From 6d48622544beeb6fb16e9c5823d9a6ab6eac1eb8 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Sun, 18 Oct 2026 22:59:53 -0700 Subject: [PATCH 01/28] [TRTLLM-17033][feat] Engine, QSA and modeling_v2 catalog support for the Qwen3.8-Flash-Next features The changes outside the target directory that its MTP, video, KV-storage and breakable prefill graph features need, with their unit tests: - thop attention op: read an NVFP4 KV cache through the FP8 context FMHA. - QSA cache manager: pack the side planes so an NVFP4 KV cache keeps one page-index scale. QSA metadata: refresh the draft table; drop the unused draft-forward saved state. QSA kernels: `displace_oldest`, so a full IndexShare row gives up its oldest ids, not its newest. - Decoder runner: carry the mRoPE position delta through MTP verify groups; refuse a first-draft request on an mRoPE model, and define the first-draft requests that check reads. - Model config: `moe_backend_requested`, the MoE backend asked for before AUTO resolves. Model loader: `effective_kv_cache_quant` and `effective_mamba_ssm_cache_dtype`, shared with the modeling_v2 router. Executor: a model's optional `validate_request` hook at admission. - modeling_v2: the router reads the effective KV-cache quantization and SSM-state dtype; the catalog's thop_attention is certified over FP8 and NVFP4 pools at the qwen4_exp cell; `eager_on_graph` runs a target's breaks eagerly under the breakable prefill CUDA graph. Folds these commits of the review branch (their changes to the files of this commit): - Run the Qwen3.8-Flash-Next trunk routed experts on CUTEDSL (5405a3d6ee) - Complete MTP for the Qwen3.8-Flash-Next modeling_v2 target (0a1b9c48da) - Accept video input in the Qwen3.8-Flash-Next modeling_v2 target (80cb4ba8b9) - Carry the mRoPE position delta through MTP verify groups (cb2bd1099b) - Store the Qwen3.8-Flash-Next KV cache as FP8 when kv_cache_quant_algo is FP8 (744607bf5c) - Decline engine configurations the qwen4_exp target cannot run (311920b003) - Serve video past one tower call and reject an oversize frame at admission (25f5127266) - Pin the mRoPE delta read slots of speculative verify groups in the decoder runner (b179f610e9) - Let a full IndexShare row give up its oldest ids, not its newest (6d498a45ad) - Drop QSA metadata's unused draft-forward saved state and test the draft table refresh (60cef34d78) - Refuse a first-draft request on an mRoPE model in the decoder runner (1b9ecd33c5) - Run the Qwen3.8-Flash-Next modeling_v2 target at MTP draft length 1..7 under GDN cached replay (bfa50da860) - Serve the Qwen3.8-Flash-Next modeling_v2 Gated-DeltaNet state in bf16 (a0f2206796) - Serve the qwen4_exp modeling_v2 target under the breakable prefill CUDA graph (0f3c3c078c) - Pack the QSA side planes so an NVFP4 KV cache keeps one page-index scale (fe6fb439bc) - Read an NVFP4 KV cache through the FP8 context FMHA in the thop attention op (bc888a99d1) - Certify attention/thop_attention over an NVFP4 KV cache at the qwen4_exp cell (2a29fcd08a) - Store the qwen4_exp KV cache as NVFP4 when kv_cache_quant_algo is NVFP4 (e3c62d38f0) - Define the first-draft requests the mRoPE refusal reads (523e5c40c6) Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- cpp/tensorrt_llm/thop/attentionOp.cpp | 4 +- .../_experimental/modeling_v2/README.md | 13 +- .../modeling_v2/_router_index.py | 49 ++- .../catalog/attention/thop_attention.py | 329 ++++++++++++++- .../modeling_v2/catalog/graph/__init__.py | 3 + .../catalog/graph/eager_on_graph.py | 111 +++++ .../modeling_v2/models/deepseek_v3/routing.py | 1 + .../modeling_v2/models/gpt_oss/routing.py | 1 + .../backends/sparse/qsa/cache_manager.py | 254 +++++++++-- .../attention/backends/sparse/qsa/kernels.py | 23 +- .../attention/backends/sparse/qsa/metadata.py | 99 ++++- tensorrt_llm/_torch/model_config.py | 3 + .../engine/runners/decoder/inputs.py | 127 +++++- .../_torch/pyexecutor/model_loader.py | 30 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 6 + .../qsa/test_qsa_cache_manager_layout.py | 224 ++++++++++ .../attention/sparse/qsa/test_qsa_sparse.py | 299 +++++++++++++ .../test_mrope_first_draft_refusal.py | 54 +++ .../_torch/executor/test_py_executor.py | 86 ++++ .../executor/test_pytorch_model_engine.py | 399 +++++++++++++++++- .../modeling/test_qsa_runtime_wiring.py | 2 +- .../test_modeling_v2_thop_attention.py | 178 +++++++- .../graph/test_modeling_v2_eager_on_graph.py | 285 +++++++++++++ .../modeling_v2/test_modeling_v2_routing.py | 199 ++++++++- 24 files changed, 2661 insertions(+), 118 deletions(-) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/catalog/graph/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/catalog/graph/eager_on_graph.py create mode 100644 tests/unittest/_torch/attention/sparse/qsa/test_qsa_cache_manager_layout.py create mode 100644 tests/unittest/_torch/executor/test_mrope_first_draft_refusal.py create mode 100644 tests/unittest/_torch/modeling_v2/graph/test_modeling_v2_eager_on_graph.py diff --git a/cpp/tensorrt_llm/thop/attentionOp.cpp b/cpp/tensorrt_llm/thop/attentionOp.cpp index 7fe0694b9b36..f588dd465a45 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.cpp +++ b/cpp/tensorrt_llm/thop/attentionOp.cpp @@ -1355,8 +1355,10 @@ void attention(torch::Tensor q, std::optional k, std::optionalmRotaryEmbeddingLongMscale = rope_long_m_scale; op->mRotaryEmbeddingMaxPositions = rope_max_positions; op->mRotaryEmbeddingOriginalMaxPositions = rope_original_max_positions; + // Only the FP8 context FMHA kernels read an NVFP4 KV cache, whatever the output type. MLA keeps its own NVFP4 KV + // handling (mUseNvfp4MlaKvCache, mFP8ContextMLA below) and must not pick this up. op->mFP8ContextFMHA = is_fp8_out || is_fp4_out || (op->mKVCacheQuantMode.hasFp8KvCache() && use_paged_context_fmha) - || use_sage_attn; + || (!is_mla_enable && op->mKVCacheQuantMode.hasFp4KvCache()) || use_sage_attn; // SageAttention block sizes and quantization mode. op->mSageAttnNumEltsPerBlkQ = static_cast(sage_attn_num_elts_per_blk_q); op->mSageAttnNumEltsPerBlkK = static_cast(sage_attn_num_elts_per_blk_k); diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md index 5e34934827bc..8528a8c20827 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md @@ -36,10 +36,13 @@ exists to catch. The decision is one function, `route`, in the family's `routing.py`. It reads the facts that fix which code runs -- the checkpoint's config shape, the -parallel mapping, the GPU's compute capability and the speculative-decoding -config -- and either returns a target class or raises `NoTarget` with the -reason in words. Tuning knobs such as `max_batch_size`, `max_num_tokens` and CUDA graph -settings are not routing inputs: they do not change which code runs, so a +parallel mapping, the GPU's compute capability, the speculative-decoding +config and the engine choices in `EngineFacts` (KV cache quantization, MoE +backend, recurrent-state cache dtype, prefill CUDA graph backend) -- and either returns a target class or +raises `NoTarget` with the reason in words. A target declines what its +constructor would refuse, so that `auto` falls back instead of failing mid-build. +Tuning knobs such as `max_batch_size`, `max_num_tokens` and the CUDA graph +bucket lists are not routing inputs: they do not change which code runs, so a target claims a configuration regardless of them. `require` therefore promises that the target ran, not that a particular knob setting was measured. The model loader decides when it is about to build the model, after the built-in model's @@ -67,7 +70,7 @@ LLM(model=..., modeling_v2=...) -> ModelLoader | _router_index.py +-- architectures[0] -> routing module models//routing.py +-- route(config, mapping, compute_capability, - | spec_config, kv_cache_quant_algo) + | spec_config, kv_cache_quant_algo, engine) | -> target class, or NoTarget | config.modeling_v2_target = class diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py index e8228918397e..c9b4c681f79c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py @@ -17,10 +17,11 @@ architecture family does. * The routing module's ``route`` owns the rest of the decision as one forward-reading tree over the facts that fix which code runs: the - checkpoint's config, the parallel mapping, the GPU's compute capability and - the speculative-decoding config. It returns the target class, or raises - :class:`NoTarget` with the reason in words. Reading that one file tells you - where any configuration lands. + checkpoint's config, the parallel mapping, the GPU's compute capability, the + speculative-decoding config, the KV-cache algorithm the engine will run with + and the engine choices in :class:`EngineFacts`. It returns the target class, + or raises :class:`NoTarget` with the reason in words. Reading that one file + tells you where any configuration lands. Tuning knobs -- ``max_batch_size``, ``max_num_tokens``, CUDA graph settings -- are deliberately not routing inputs. They do not change which code runs, so a @@ -37,7 +38,7 @@ import enum import os from importlib import import_module -from typing import TYPE_CHECKING, Optional, Tuple +from typing import TYPE_CHECKING, Any, NamedTuple, Optional, Tuple from tensorrt_llm.logger import logger @@ -165,6 +166,43 @@ def _kv_cache_quant_algo(config: "ModelConfig", llm_args: "TorchLlmArgs"): return {"fp8": QuantAlgo.FP8, "nvfp4": QuantAlgo.NVFP4}.get(dtype, dtype) +class EngineFacts(NamedTuple): + """Engine choices that fix which code runs but are not in the checkpoint's config. + + A target's constructor refuses what it cannot run, but by then ``route`` has + already claimed the deployment and ``auto`` can no longer fall back; a route + that reads these declines first. + """ + + torch_dtype: Any # torch.dtype the engine runs the model in + kv_cache_quant_algo: Any # QuantAlgo or None: the checkpoint's under "auto", else the override + moe_backend: str # `moe_config.backend` as requested, before AUTO resolves + mamba_ssm_cache_dtype: Any # torch.dtype of the recurrent-state cache + gdn_replay: bool # TRTLLM_USE_GDN_REPLAY, which the engine reads for eligible MTP workloads + prefill_cuda_graph_backend: str # `prefill_cuda_graph_backend`: disabled, piecewise, breakable + + +def engine_facts(config: "ModelConfig", llm_args: "TorchLlmArgs") -> EngineFacts: + # Imported here: the model loader imports this module. + from tensorrt_llm._torch.pyexecutor.model_loader import ( + effective_kv_cache_quant, + effective_mamba_ssm_cache_dtype, + ) + from tensorrt_llm._torch.utils import is_gdn_replay_enabled + + kv_cache_config = llm_args.kv_cache_config + return EngineFacts( + torch_dtype=config.torch_dtype, + kv_cache_quant_algo=effective_kv_cache_quant(config, kv_cache_config.dtype), + moe_backend=config.moe_backend_requested, + mamba_ssm_cache_dtype=effective_mamba_ssm_cache_dtype( + config, kv_cache_config.mamba_ssm_cache_dtype + ), + gdn_replay=is_gdn_replay_enabled(), + prefill_cuda_graph_backend=str(llm_args.prefill_cuda_graph_backend.value), + ) + + def modeling_v2_resolve(config: "ModelConfig", llm_args: "TorchLlmArgs") -> Optional[type]: """Decide which modeling_v2 target, if any, this deployment builds. @@ -195,6 +233,7 @@ def modeling_v2_resolve(config: "ModelConfig", llm_args: "TorchLlmArgs") -> Opti _device_compute_capability(), config.spec_config, _kv_cache_quant_algo(config, llm_args), + engine=engine_facts(config, llm_args), ) except NoTarget as why: if mode is ModelingV2Mode.REQUIRE: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/thop_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/thop_attention.py index fed8d5fa4755..80a2ad211f21 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/thop_attention.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/thop_attention.py @@ -26,6 +26,10 @@ #: quant_mode bit for an fp8 KV cache. QM_FP8_KV_CACHE = 128 +#: quant_mode bit for an NVFP4 KV cache: e2m1 values in the pool, one e4m3 scale +#: per 16 values in a second pool, fp32 [3] kv scales (q, K, V). +QM_NVFP4_KV_CACHE = 8192 + #: qwen4_exp's full-attention layer -- the twelve trunk layers and the MTP #: module alike: packed-QKV causal GQA over a bf16 pool, rope applied outside, #: the paged-context path on (the target serves with block reuse and chunked @@ -63,6 +67,31 @@ mask_type=0, ) +#: The qwen4_exp full-attention call over an fp8-e4m3 pool at the production kv +#: scale 1.0 (fp32 [1] scale tensors, quant_mode 128): what the trunk layers and +#: the MTP draft layer pass when the engine serves an fp8 KV cache. Same chosen +#: pool sizes as `_QWEN4_DENSE`. +_QWEN4_FP8 = dict(_QWEN4_DENSE, quant_mode=QM_FP8_KV_CACHE, kv_scale=1.0) + +#: The same call over an NVFP4 pool: e2m1 data and a second pool of e4m3 block +#: scales, fp32 [3] unit kv scales, `[1, 2, 2]` pool pointers, bf16 output -- +#: what the trunk layers and the MTP draft layer pass when the engine serves an +#: NVFP4 KV cache. Same chosen pool sizes as `_QWEN4_DENSE`. +_QWEN4_NVFP4 = dict(_QWEN4_DENSE, quant_mode=QM_NVFP4_KV_CACHE, kv_scale=1.0) + +#: `_QWEN4_DENSE` and `_QWEN4_FP8` for a verify step of up to 32 requests, the captured-batch +#: ceiling: one block-offset row per request and pages for 32 histories (the largest cell +#: needs 133 of the 144). +_QWEN4_DENSE_G32 = dict(_QWEN4_DENSE, max_num_requests=32, num_blocks=144) +_QWEN4_FP8_G32 = dict(_QWEN4_FP8, max_num_requests=32, num_blocks=144) + + +def _spread(batch: int, longest: int) -> tuple[int, ...]: + """`batch` histories rising evenly from 1 to `longest` tokens: one-page, page-boundary + and multi-page sequences in one verify step.""" + return tuple(1 + (longest - 1) * i // (batch - 1) for i in range(batch)) + + #: A GQA 24/4/256 full-attention layer (Qwen3.8-27B) over an fp8-e4m3 pool at #: the production kv scale 1.0. Same chosen pool sizes as `_QWEN4_DENSE`. _FP8_24_4_256 = dict( @@ -90,8 +119,9 @@ class ThopAttention(OpWrapper): it, by different paths: gpt-oss on the standard path with attention sinks and an alternating sliding window, deepseek on MLA's two phases, and qwen4_exp on the standard path twice over -- its full-attention layers and - MTP module (GQA 24/2/256, a P-token draft block per generation sequence, - a two-pool hybrid mapping), and its vision tower with no KV cache at all. + MTP module (GQA 24/2/256 over a bf16 or an fp8-e4m3 pool, a P-token draft + block per generation sequence, a two-pool hybrid mapping), and its vision + tower with no KV cache at all. Writes rows `[:num_tokens]` of `output` and returns `output` -- the same buffer, mutated and handed back so a call site can show what it touched @@ -277,7 +307,7 @@ class ThopAttention(OpWrapper): ), ), Cell( - why="qwen4_exp MTP verify P=4 (max_draft_len 3): the target's ceiling", + why="qwen4_exp MTP verify P=4 (max_draft_len 3)", spec=dict( _QWEN4_DENSE, cached=(70, 35), @@ -286,6 +316,61 @@ class ThopAttention(OpWrapper): predicted_tokens_per_seq=4, ), ), + *( + Cell( + why=f"qwen4_exp MTP verify P={p} (max_draft_len {p - 1})" + + (": the target's ceiling" if p == 8 else ""), + spec=dict( + _QWEN4_DENSE, + cached=(70, 35), + new=(p, p), + num_contexts=0, + predicted_tokens_per_seq=p, + ), + ) + for p in (5, 6, 7, 8) + ), + Cell( + why="qwen4_exp P=8 G=1: one draft block from row 95, across the page edge at 96", + spec=dict( + _QWEN4_DENSE, + cached=(95,), + new=(8,), + num_contexts=0, + predicted_tokens_per_seq=8, + ), + ), + # The trtllm-gen launch picks its generation kernel by batch size at P 5..8, so each P + # is also driven at 32 requests and at one batch in 17..20, with histories of every + # page shape in the step. + *( + Cell( + why=f"qwen4_exp bf16 batch P={p} G=32: 32 blocks over histories 1..218, " + f"{32 * p} q rows", + spec=dict( + _QWEN4_DENSE_G32, + cached=_spread(32, 218), + new=(p,) * 32, + num_contexts=0, + predicted_tokens_per_seq=p, + ), + ) + for p in (5, 6, 7, 8) + ), + *( + Cell( + why=f"qwen4_exp bf16 batch P={p} G={g}: {g} blocks over histories 200..1, " + f"{g * p} q rows", + spec=dict( + _QWEN4_DENSE_G32, + cached=tuple(reversed(_spread(g, 200))), + new=(p,) * g, + num_contexts=0, + predicted_tokens_per_seq=p, + ), + ) + for p, g in ((5, 17), (6, 18), (7, 19), (8, 20)) + ), Cell( why="qwen4_exp mixed step: a fresh 20-row context ahead of two P=2 " "draft blocks, generation rows token-major after the context rows", @@ -297,6 +382,17 @@ class ThopAttention(OpWrapper): predicted_tokens_per_seq=2, ), ), + Cell( + why="qwen4_exp mixed step P=8: a fresh 20-row context ahead of two P=8 " + "draft blocks over histories 70 and 35", + spec=dict( + _QWEN4_DENSE, + cached=(0, 70, 35), + new=(20, 8, 8), + num_contexts=1, + predicted_tokens_per_seq=8, + ), + ), Cell( why="qwen4_exp hybrid pool: row 47 of the 48-row two-pool mapping " "(pool 0, layer-in-pool 11), P=2 verify", @@ -310,6 +406,19 @@ class ThopAttention(OpWrapper): predicted_tokens_per_seq=2, ), ), + Cell( + why="qwen4_exp P=8 hybrid pool: row 47 of the 48-row two-pool mapping " + "(pool 0, layer-in-pool 11), P=8 verify", + spec=dict( + _QWEN4_DENSE, + pool_layout="hybrid48", + local_layer_idx=47, + cached=(64, 25), + new=(8, 8), + num_contexts=0, + predicted_tokens_per_seq=8, + ), + ), Cell( why="48-layer single pool, identity mapping: row 47 = layer-in-pool 47, " "24 rows over a 40-token cached prefix", @@ -371,6 +480,24 @@ class ThopAttention(OpWrapper): "12544-row frame is the longest segment", spec=dict(_QWEN4_VISION, seg_lens=(12544, 9600)), ), + # A vision-tower call is one frame per frame group, at most 32768 rows per call. + Cell( + why="video: 384 frame groups of 64 rows (8x8 patches), the untiled " + "kernel's longest frame", + spec=dict(_QWEN4_VISION, seg_lens=(64,) * 384), + ), + Cell( + why="video: 256 frame groups of 128 rows, the 32768-row ceiling", + spec=dict(_QWEN4_VISION, seg_lens=(128,) * 256), + ), + Cell( + why="video: 2 frame groups of 16384 rows (128x128), the longest frame at the ceiling", + spec=dict(_QWEN4_VISION, seg_lens=(16384, 16384)), + ), + Cell( + why="video: 8192 frame groups of 2x2 patches, the most a 32768-row call can carry", + spec=dict(_QWEN4_VISION, seg_lens=(4,) * 8192), + ), Cell( why="fp8 pool packed prefill, GQA 24q/4kv d256: the FMHA reads the " "bf16 q rows, not the pool", @@ -479,6 +606,165 @@ class ThopAttention(OpWrapper): predicted_tokens_per_seq=3, ), ), + # qwen4_exp's own geometry over an fp8 pool: the kv head count changes the compiled + # decode kernel, so the 24/4/256 cells above do not cover it. + Cell( + why="qwen4_exp fp8 pool paged-context prefill, 24/2/256: a fresh 48-row " + "sequence beside 33 rows over a 48-token e4m3 prefix", + spec=dict( + _QWEN4_FP8, + cached=(0, 48), + new=(48, 33), + num_contexts=2, + predicted_tokens_per_seq=1, + ), + ), + Cell( + why="qwen4_exp fp8 pool decode P=1, 24/2/256: histories 200 and 48 span " + "pages and KV tiles", + spec=dict( + _QWEN4_FP8, + cached=(200, 48), + new=(1, 1), + num_contexts=0, + predicted_tokens_per_seq=1, + ), + ), + *( + Cell( + why=f"qwen4_exp fp8 pool MTP verify P={p}, 24/2/256 (max_draft_len {p - 1})", + spec=dict( + _QWEN4_FP8, + cached=(70, 35), + new=(p, p), + num_contexts=0, + predicted_tokens_per_seq=p, + ), + ) + for p in (2, 4, 5, 6, 7, 8) + ), + Cell( + why="qwen4_exp fp8 pool mixed step P=3, 24/2/256: a 20-row context ahead " + "of two P=3 blocks over histories 31 and 64", + spec=dict( + _QWEN4_FP8, + cached=(0, 31, 64), + new=(20, 3, 3), + num_contexts=1, + predicted_tokens_per_seq=3, + ), + ), + Cell( + why="qwen4_exp fp8 pool mixed step P=8, 24/2/256: a 20-row context ahead " + "of two P=8 blocks over histories 31 and 64", + spec=dict( + _QWEN4_FP8, + cached=(0, 31, 64), + new=(20, 8, 8), + num_contexts=1, + predicted_tokens_per_seq=8, + ), + ), + # KV stays within one 128-key tile (histories to 119 plus 8 rows), where the e4m3 + # probability model of `reference` is the kernel's own at any batch. + Cell( + why="qwen4_exp fp8 batch P=8 G=32: 32 blocks over histories 1..119, 256 q rows", + spec=dict( + _QWEN4_FP8_G32, + cached=_spread(32, 119), + new=(8,) * 32, + num_contexts=0, + predicted_tokens_per_seq=8, + ), + ), + Cell( + why="qwen4_exp fp8 hybrid pool P=2: row 47 of the 48-row two-pool mapping " + "(pool 0, layer-in-pool 11) over e4m3 slabs", + spec=dict( + _QWEN4_FP8, + pool_layout="hybrid48", + local_layer_idx=47, + cached=(64, 25), + new=(2, 2), + num_contexts=0, + predicted_tokens_per_seq=2, + ), + ), + Cell( + why="qwen4_exp fp8 hybrid pool P=8: row 47 of the 48-row two-pool mapping " + "(pool 0, layer-in-pool 11) over e4m3 slabs", + spec=dict( + _QWEN4_FP8, + pool_layout="hybrid48", + local_layer_idx=47, + cached=(64, 25), + new=(8, 8), + num_contexts=0, + predicted_tokens_per_seq=8, + ), + ), + Cell( + why="qwen4_exp nvfp4 pool paged-context prefill, 24/2/256: a fresh 48-row " + "sequence beside 33 rows over a 48-token e2m1 prefix", + spec=dict( + _QWEN4_NVFP4, + cached=(0, 48), + new=(48, 33), + num_contexts=2, + predicted_tokens_per_seq=1, + ), + ), + Cell( + why="qwen4_exp nvfp4 pool decode P=1, 24/2/256: histories 200 and 48", + spec=dict( + _QWEN4_NVFP4, + cached=(200, 48), + new=(1, 1), + num_contexts=0, + predicted_tokens_per_seq=1, + ), + ), + Cell( + why="qwen4_exp nvfp4 pool MTP verify P=2, 24/2/256 (max_draft_len 1)", + spec=dict( + _QWEN4_NVFP4, + cached=(70, 35), + new=(2, 2), + num_contexts=0, + predicted_tokens_per_seq=2, + ), + ), + Cell( + why="qwen4_exp nvfp4 pool MTP verify P=4, 24/2/256 (max_draft_len 3)", + spec=dict( + _QWEN4_NVFP4, + cached=(70, 35), + new=(4, 4), + num_contexts=0, + predicted_tokens_per_seq=4, + ), + ), + Cell( + why="qwen4_exp nvfp4 pool mixed step P=1: a 20-row context ahead of two " + "1-row sequences over histories 31 and 64", + spec=dict( + _QWEN4_NVFP4, + cached=(0, 31, 64), + new=(20, 1, 1), + num_contexts=1, + predicted_tokens_per_seq=1, + ), + ), + Cell( + why="qwen4_exp nvfp4 pool mixed step P=4: a 20-row context ahead of two 4-row blocks", + spec=dict( + _QWEN4_NVFP4, + cached=(0, 31, 64), + new=(20, 4, 4), + num_contexts=1, + predicted_tokens_per_seq=4, + ), + ), ) note = """ @@ -510,24 +796,51 @@ class its own layer group, so its mapping carries two pool ids; the op reads `is_spec_decoding_enabled` False, which is what sm_103 passes for a linear draft, and mask_type 0 and 1 are bitwise identical at P > 1. Nothing checks q's row count against P. P reaches the decode kernel as `maxSeqLenQ`: one - JIT compile per P. + JIT compile per P. P is driven at 1..8 (max_draft_len 1..7) at 24/2/256 over + bf16 and fp8 pools. From P = 5 the launch picks its generation kernel by + batch size, so P 5..8 is also driven at 32 requests and at one batch in + 17..20 (bf16), and P = 8 at 32 requests over the fp8 pool. No KV cache (an encoder): pass `kv_cache_block_offsets`, `host_kv_cache_pool_pointers` and `host_kv_cache_pool_mapping` all None, with `update_kv_cache=True` -- False is rejected ("KV cache update cannot be disabled now"). One context sequence per frame; mask_type 0 confines a frame to itself. `tokens_per_block` and `use_paged_context_fmha` are - bitwise inert there. Passing only some of the three None is unmeasured. + bitwise inert there. Passing only some of the three None is unmeasured. An + image is one frame and a video one frame per frame group; the cells run up to + 8192 frames and 32768 rows per call at 16/16/72 under the padding mask. + `FusedMHARunnerV2` picks a different fmha_v2 kernel for a call whose longest + frame is at most 64 rows than for a longer one, so a frame's last bits can + move with the other frames it shares a call with, within `reference`'s band. head_size 72 exists in the context phase only: a generation call raises "unsupported head_size 72", and the process stays usable. Over an fp8-e4m3 pool the packed context path reads the bf16 q rows and never the pool, while `use_paged_context_fmha=True` reads the pool, so its operands are e4m3: the flag is output-inert over a bf16 pool only. Every - fp8 cell runs kv scale 1.0; at 1.5 `reference` does not reproduce the - kernel's probability rounding. + fp8 cell runs kv scale 1.0 through fp32 [1] scale tensors; at 1.5 `reference` + does not reproduce the kernel's probability rounding. quant_mode 384, 1152, 4224 and 4480 are bit-identical to 128 at 24/4/256; - 4352 (the same bits without FP8_KV_CACHE) reads e4m3 slabs as bf16. + 4352 (the same bits without FP8_KV_CACHE) reads e4m3 slabs as bf16. At + 24/2/256 only 128 was driven, with KV up to 200 keys at P = 1 and up to + 127 at P 2..8: the fp8 kernels round the softmax probabilities against a + per-tile max, which the row-max model of + `reference(e4m3_probabilities=True)` reproduces within one 128-key tile + and approximates up to 200 keys. + + NVFP4 pool (quant_mode 8192): K and V are stored as e2m1 values, two per + byte (low nibble = even element), per head page `[tpb, D/2]`, with one e4m3 + scale per 16 values in a second pool whose pointer is the third axis of + `host_kv_cache_pool_pointers` (`[pools, 2, 2]`: primary/secondary x + data/scales); K's scales are token-major per head page, V's in four-token + groups `[tpb/4, D/16, 4]`. The append quantizes with `kv_scale_orig_quant[1]` + / `[2]` as second-level scales (fp32 `[3]`, `[0]` is q's e4m3 scale). One + block-offset table addresses both pools, so they must hold the same pages + per slot. The kernels are trtllm-gen cubins with e4m3 q and e2m1 K/V and + exist for the paged-context path only, which needs the op to enable FP8 + context FMHA for an FP4 KV cache. Every NVFP4 cell runs the paged-context + path at kv scale 1.0 and reads K and V through e4m3, so `reference` runs + over the NVFP4 values with `e4m3_inputs` and `e4m3_probabilities`. The sparse / skip-softmax arguments are bitwise inert in the standard configuration (nine off-default settings at 24/2/256): passing them does diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/graph/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/graph/__init__.py new file mode 100644 index 000000000000..c976f03c8914 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/graph/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Graph-capture entries.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/graph/eager_on_graph.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/graph/eager_on_graph.py new file mode 100644 index 000000000000..43f18a5ed022 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/graph/eager_on_graph.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One break of a breakable CUDA graph: a callable run eagerly between captured segments. + +Wraps the engine's `eager_on_graph` (`pyexecutor/breakable_cuda_graph`). It launches no +kernel of its own; what it certifies is the control flow around the callable it is handed. +""" + +from __future__ import annotations + +from typing import Any, Callable + +import torch + +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import eager_on_graph as _eager_on_graph + +from .._op import Arch, Cell, OpWrapper + + +class EagerOnGraph(OpWrapper): + """`inner(*args)`, eagerly; under breakable CUDA-graph capture, a break. + + Outside a `BreakableCUDAGraphCapture` the call is a passthrough. Inside one it ends the + captured segment, runs `inner(*args)` once, records a replay function holding `inner` + and the arguments, and opens the next segment. `BreakableCUDAGraph.replay` then runs + each segment and, between two of them, re-runs `inner` with the **capture-time** + arguments: a CUDA tensor by its capture-time address, a CPU tensor as the same object, + an `int` / `float` / `bool` / `str` / `None` frozen at its capture-time value. Whatever + varies per step therefore reaches `inner` through state it reads itself, which the + caller refreshes before each replay, never through an argument. + + Fusion boundary: none. The break runs exactly what `inner` runs, on the current stream. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # Spec axes: `capture` (inside a breakable capture, or plain eager), `breaks` (breaks in + # one body), `tokens` (the bucket `T` the body is captured at), `width` (row width), + # `live_rows` (the rows `n` each replay carries, set together with a fresh per-step scale + # and fresh input rows). Each break's `inner` reads `n` and the scale from live state, + # mixes rows `[:n]` across the batch, writes them into a fixed in-place output, zeroes + # rows `[n:]` and returns None: the shape of every break a prefill body makes. + CELLS: tuple[Cell, ...] = ( + Cell( + why="outside capture: a passthrough, bitwise the inner call", + spec=dict(capture=False, breaks=2, tokens=256, width=128, live_rows=(256, 17)), + ), + Cell( + why="one break, bucket 256: replays at new live rows equal eager", + spec=dict(capture=True, breaks=1, tokens=256, width=128, live_rows=(17, 256, 1)), + ), + Cell( + why="50 breaks in one body, bucket 1024", + spec=dict(capture=True, breaks=50, tokens=1024, width=128, live_rows=(1000, 255, 1024)), + ), + Cell( + why="bucket 8192, live rows across the 2048 boundary and the full bucket", + spec=dict(capture=True, breaks=2, tokens=8192, width=64, live_rows=(6144, 8192, 2047)), + ), + ) + + note = """ + Preconditions and hazards, each silent unless stated otherwise: + + * **Scalars are frozen.** An `int`, `float`, `bool`, `str` or `None` argument replays at + its capture-time value. A static per-site constant (a layer index) is right; a + per-step count is wrong at every replay. + * **CUDA tensors are re-read by address.** An argument must live at a stable address for + the graph's life: allocated inside the captured body (the graph pool) or an + engine-owned persistent buffer. A tensor allocated per step outside the graph replays + the capture-time memory. + * **Any other argument raises.** A dataclass, a namespace or a bound object raises + TypeError at capture, after `inner` already ran once. Loud, but it is a break raising + under capture, so the next hazard follows it. + * **A break that raises under capture leaves a CUDA error pending.** The raise comes + after the break ended the captured segment and before it opened the next, and the + capture's exit ends the ended segment again: it raises RuntimeError over the break's + own error and leaves cudaErrorIllegalState pending in the CUDA runtime, so the + process's next CUDA launch fails once with "the operation cannot be performed in the + present state". A caller that recovers must take the pending error + (`cudaGetLastError`) before launching anything else. + * **Side-stream waits are dropped.** While a capture is open, including inside `inner`'s + own capture-time run, a wait of the capture stream on a stream that is not capturing + is dropped with a warning, and replay does not preserve it. + * **The return value is copied each replay** onto the object `inner` returned at + capture. `inner` writing a fixed output in place returns None. + * **A fixed output must be fully written.** Rows `inner` leaves alone keep whatever an + earlier replay or another bucket left there. + + Hard limits + none: no index-typed argument + """ + + def __call__(self, *args: Any, inner: Callable[..., Any]) -> Any: + return _eager_on_graph(inner)(*args) + + def reference(self, *args: Any, inner: Callable[..., Any]) -> Any: + """The inner call, eagerly: what every replay of the break must reproduce.""" + return inner(*args) + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """Bitwise: a break adds no arithmetic, so any difference is a frozen or stale + operand, not rounding.""" + assert out.dtype == ref.dtype, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + assert torch.equal(out, ref), ( + f"{int((out != ref).sum())} of {out.numel()} elements differ from the eager body" + ) + + +eager_on_graph = EagerOnGraph() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/routing.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/routing.py index 6c2ebcaf4a15..8d0502f23bcc 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/routing.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/routing.py @@ -76,6 +76,7 @@ def route( compute_capability: Tuple[int, int], spec_config: Optional[Any], kv_cache_quant_algo: Optional[Any], + engine: Optional[Any] = None, ) -> type: """Return the target class for this deployment, or raise ``NoTarget`` saying why none fits. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/routing.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/routing.py index 3582cc17a695..e18cc12e1644 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/routing.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/routing.py @@ -49,6 +49,7 @@ def route( compute_capability: Tuple[int, int], spec_config: Optional[Any], kv_cache_quant_algo: Optional[Any], + engine: Optional[Any] = None, ) -> type: """Return the target class for this deployment, or raise ``NoTarget`` saying why none fits. diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py index a5ddcaa70614..442cb8d67758 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py @@ -11,6 +11,7 @@ from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import Role from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 from tensorrt_llm._utils import TensorWrapper, binding_to_torch_dtype, convert_to_torch_tensor +from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp from tensorrt_llm.runtime.kv_cache_manager_v2 import BufferConfig, DataRole, PageIndexMode @@ -31,10 +32,22 @@ _SHARED_POSITION_ROLE_INDEX = 0 _INDEX_K_ELEMENT_BYTES = torch.empty((), dtype=QSA_INDEX_K_CACHE_DTYPE).element_size() _POSITION_ELEMENT_BYTES = torch.empty((), dtype=QSA_POSITION_CACHE_DTYPE).element_size() +# NVFP4 packs two e2m1 values per byte and keeps one e4m3 scale byte per +# 16-element block. +_NVFP4_VALUES_PER_BYTE = 2 +_NVFP4_BLOCK_SIZE = 16 +# Granularity by which the packed side-plane page grows until its size is +# distinct from every main-cache page; keeps each plane 16-byte aligned. +_SIDE_PLANES_SIZE_STEP = 128 # Per-token RoPE/mRoPE coordinates used when raw index keys are compressed. # They are request state shared by all local QSA layers. QSA_INDEX_POSITION = DataRole("qsa_index_position") +# Every local layer's index-K page followed by the position page, as one +# buffer per lifecycle slot. Registered instead of the per-layer INDEX_KEY and +# the position role when either page would be the size of a K/V or +# block-scale page (see `_extra_buffers_per_layer`). +QSA_SIDE_PLANES = DataRole("qsa_side_planes") class QSAMambaHybridCacheManagerV2(MambaHybridCacheManagerV2): @@ -47,6 +60,14 @@ class QSAMambaHybridCacheManagerV2(MambaHybridCacheManagerV2): Index geometry is resolved from the same serving/checkpoint configuration as the indexer. A local fallback could otherwise give the indexer and its paged cache different shapes or page strides. + + V2 storage coalesces the buffers of one lifecycle by page size alone. The + attention op addresses a layer's K/V pages and its NVFP4 block-scale pages + with one block-offset table, so a side page that lands in the K/V pool + (NVFP4 at 32 tokens per block: index-K 8192 B == K == V) changes the K/V + slot stride but not the block-scale one. The side planes are then + registered as one packed per-slot buffer whose size matches no main-cache + page, and this manager serves their per-layer views itself. """ def __init__( @@ -76,6 +97,9 @@ def __init__( layer_idx for layer_idx, enabled in enumerate(layer_mask) if enabled ] self.qsa_position_layer_id: Optional[int] = None + # Bytes of one packed side-plane page, or None while the side planes + # are registered as per-layer roles. Set by `_extra_buffers_per_layer`. + self.qsa_side_planes_size: Optional[int] = None super().__init__(*args, layer_mask=layer_mask, **kwargs) def get_disagg_role_mapper_kinds(self) -> Dict[DataRole, MapperKind]: @@ -86,6 +110,7 @@ def get_disagg_role_mapper_kinds(self) -> Dict[DataRole, MapperKind]: return { **super().get_disagg_role_mapper_kinds(), QSA_INDEX_POSITION: MapperKind.REPLICATED, + QSA_SIDE_PLANES: MapperKind.REPLICATED, } @cached_property @@ -110,37 +135,106 @@ def _extra_buffers_per_layer( *, tokens_per_block: int, ) -> dict[int, list[BufferConfig]]: - """Register per-layer index K and lifecycle-aligned position pages.""" + """Register per-layer index K and lifecycle-aligned position pages. + + Both are registered as their own roles unless either page has the + size of a K/V or block-scale page, which V2 would coalesce into that + pool. They are then packed, in local sparse-layer order with the + position page last, into one `QSA_SIDE_PLANES` buffer on the first + local sparse layer, grown until its size matches no main-cache page. + """ index_size = ( self.qsa_index_kv_heads * self.qsa_index_dim * _INDEX_K_ELEMENT_BYTES * tokens_per_block ) + position_size = QSA_POSITION_COORDINATE_AXES * _POSITION_ELEMENT_BYTES * tokens_per_block local_sparse_layers = self.qsa_local_sparse_layer_ids # Coordinates are request-wide, so all local indexers use one view. # Register the position role once; duplicating it on every sparse layer # wastes one three-axis int32 page per layer without adding state. self.qsa_position_layer_id = next(iter(local_sparse_layers), None) + self.qsa_side_planes_size = None + if self.qsa_position_layer_id is None: + return {} + position_local_idx = self.layer_offsets[self.qsa_position_layer_id] + main_page_sizes = self._main_cache_page_sizes(tokens_per_block) + if index_size in main_page_sizes or position_size in main_page_sizes: + packed_size = len(local_sparse_layers) * index_size + position_size + while packed_size in main_page_sizes: + packed_size += _SIDE_PLANES_SIZE_STEP + self.qsa_side_planes_size = packed_size + return {position_local_idx: [BufferConfig(role=QSA_SIDE_PLANES, size=packed_size)]} result = { self.layer_offsets[layer_id]: [BufferConfig(role=Role.INDEX_KEY, size=index_size)] for layer_id in local_sparse_layers } - if self.qsa_position_layer_id is not None: - local_idx = self.layer_offsets[self.qsa_position_layer_id] - result[local_idx].append( - BufferConfig( - role=QSA_INDEX_POSITION, - size=( - QSA_POSITION_COORDINATE_AXES * _POSITION_ELEMENT_BYTES * tokens_per_block - ), - ) - ) + result[position_local_idx].append(BufferConfig(role=QSA_INDEX_POSITION, size=position_size)) return result + def _main_cache_page_sizes(self, tokens_per_block: int) -> set[int]: + """Page sizes of every local K/V and block-scale buffer, in bytes.""" + sizes = set() + for local_idx in range(self.num_local_layers): + for role in self._get_buffer_roles_for_layer(local_idx): + size = self.get_layer_bytes_per_token(local_idx, role) * tokens_per_block + if size: + sizes.add(size) + return sizes + + def _side_planes(self) -> torch.Tensor: + """uint8 ``[slots, qsa_side_planes_size]`` over the packed side-plane pool.""" + local_idx = self.layer_offsets[self.qsa_position_layer_id] + page_stride = self.impl.get_page_stride(local_idx, QSA_SIDE_PLANES) + if page_stride != self.qsa_side_planes_size: + raise RuntimeError( + "QSA side-plane page stride mismatch: " + f"expected {self.qsa_side_planes_size}, got {page_stride}" + ) + converter = self.impl.get_page_index_converter(local_idx, QSA_SIDE_PLANES) + if ( + int(converter.scale) != 1 + or int(converter.layer_offset) != 0 + or int(converter.expansion) != _IDENTITY_PAGE_INDEX_EXPANSION + ): + raise RuntimeError( + "QSA side planes share their V2 pool with another buffer: " + f"scale={int(converter.scale)}, layer_offset={int(converter.layer_offset)}, " + f"expansion={int(converter.expansion)}" + ) + addr = self.impl.get_mem_pool_base_address(local_idx, QSA_SIDE_PLANES, PageIndexMode.SHARED) + num_slots = self.impl.get_page_index_upper_bound(local_idx, QSA_SIDE_PLANES) + return convert_to_torch_tensor( + TensorWrapper(addr, torch.uint8, [num_slots, self.qsa_side_planes_size]) + ) + def get_index_k_buffer( self, layer_idx: int, kv_layout: str = "NHD", ) -> Optional[torch.Tensor]: - """Return the index-K view using the indexer's resolved geometry.""" + """Return the index-K view using the indexer's resolved geometry. + + ``[slots, tokens_per_block, heads, dim]`` (NHD) or ``[slots, heads, + tokens_per_block, dim]`` (HND), indexed by lifecycle slot. + """ + if self.qsa_side_planes_size is not None: + if kv_layout not in ("NHD", "HND"): + raise ValueError(f"Unsupported kv_layout: {kv_layout}") + if layer_idx not in self.qsa_local_sparse_layer_ids: + return None + ordinal = self.qsa_local_sparse_layer_ids.index(layer_idx) + index_size = ( + self.qsa_index_kv_heads + * self.qsa_index_dim + * _INDEX_K_ELEMENT_BYTES + * self.tokens_per_block + ) + planes = self._side_planes() + plane = planes[:, ordinal * index_size : (ordinal + 1) * index_size] + if kv_layout == "NHD": + page_shape = [self.tokens_per_block, self.qsa_index_kv_heads, self.qsa_index_dim] + else: + page_shape = [self.qsa_index_kv_heads, self.tokens_per_block, self.qsa_index_dim] + return plane.view(QSA_INDEX_K_CACHE_DTYPE).view(planes.shape[0], *page_shape) return super().get_index_k_buffer( layer_idx, num_heads=self.qsa_index_kv_heads, @@ -158,11 +252,62 @@ def get_buffers( QSA tables store lifecycle slot IDs. Preserving V2's coalesced stride is therefore required even when several layers share one physical pool. + + ``[slots, 2, heads, tokens_per_block, width]`` (HND) or ``[slots, 2, + tokens_per_block, heads, width]`` (NHD), K then V. Half, bf16 and fp8 + pools are viewed in their element type at ``width = head_dim``. An + NVFP4 pool is viewed as uint8 at ``width = head_dim // 2``: two e2m1 + values per byte, element ``2i`` in the low nibble; its block scales are + `get_block_scale_buffers`. Any other format is the parent's layout. """ - if self.dtype not in QSA_SPARSE_KV_CACHE_DTYPES: - # The parent owns packed data and scale-page layouts such as - # NVFP4. QSA returns to the regular backend for these formats. - return super().get_buffers(layer_idx, kv_layout) + if self.dtype in QSA_SPARSE_KV_CACHE_DTYPES: + return self._paired_role_view( + layer_idx, + (Role.KEY, Role.VALUE), + binding_to_torch_dtype(self.dtype), + 1, + kv_layout, + ) + if self.dtype == DataType.NVFP4: + return self._paired_role_view( + layer_idx, (Role.KEY, Role.VALUE), torch.uint8, _NVFP4_VALUES_PER_BYTE, kv_layout + ) + return super().get_buffers(layer_idx, kv_layout) + + def get_block_scale_buffers( + self, + layer_idx: int, + kv_layout: str = "NHD", + ) -> Optional[torch.Tensor]: + """Return an NVFP4 pool's e4m3 block scales for this layer, K then V. + + uint8 ``[slots, 2, heads, tokens_per_block, head_dim // 16]`` (HND) or + ``[slots, 2, tokens_per_block, heads, head_dim // 16]`` (NHD) with the + same slot indexing as `get_buffers`: one scale byte per 16 values. The + attention op writes K's scales token-major within a head page, and V's + in groups of four tokens, ``[tokens_per_block / 4, head_dim / 16, 4]`` + per head; the view is the page as bytes and leaves that order to the + reader. None for a layer this rank does not hold. + """ + if self.dtype != DataType.NVFP4: + raise ValueError(f"block scales exist only for an NVFP4 pool, not {self.dtype}") + return self._paired_role_view( + layer_idx, + (Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE), + torch.uint8, + _NVFP4_BLOCK_SIZE, + kv_layout, + ) + + def _paired_role_view( + self, + layer_idx: int, + roles: tuple[DataRole, DataRole], + dtype: torch.dtype, + values_per_element: int, + kv_layout: str, + ) -> Optional[torch.Tensor]: + """Slot-indexed view over a layer's adjacent K-side and V-side pages.""" if kv_layout not in ("NHD", "HND"): raise ValueError(f"Unsupported kv_layout: {kv_layout}") if layer_idx not in self.layer_offsets: @@ -170,21 +315,22 @@ def get_buffers( if self.kv_cache_type == CacheTypeCpp.SELFKONLY: raise NotImplementedError("QSA sparse attention requires both K and V cache buffers") + key_role, value_role = roles layer_offset = self.layer_offsets[layer_idx] - addr_key = self.impl.get_mem_pool_base_address(layer_offset, Role.KEY, PageIndexMode.SHARED) + addr_key = self.impl.get_mem_pool_base_address(layer_offset, key_role, PageIndexMode.SHARED) addr_value = self.impl.get_mem_pool_base_address( - layer_offset, Role.VALUE, PageIndexMode.SHARED + layer_offset, value_role, PageIndexMode.SHARED ) - key_stride = self.impl.get_page_stride(layer_offset, Role.KEY) - value_stride = self.impl.get_page_stride(layer_offset, Role.VALUE) + key_stride = self.impl.get_page_stride(layer_offset, key_role) + value_stride = self.impl.get_page_stride(layer_offset, value_role) if addr_key + key_stride != addr_value or key_stride != value_stride: raise RuntimeError( "QSA K/V buffers are not adjacent equal-sized V2 roles: " - f"layer={layer_idx}, key_stride={key_stride}, " - f"value_stride={value_stride}" + f"layer={layer_idx}, roles={key_role}/{value_role}, " + f"key_stride={key_stride}, value_stride={value_stride}" ) - converter = self.impl.get_page_index_converter(layer_offset, Role.KEY) + converter = self.impl.get_page_index_converter(layer_offset, key_role) scale = int(converter.scale) if scale < QSA_MAIN_KV_ROLES: raise RuntimeError( @@ -193,7 +339,7 @@ def get_buffers( if int(converter.expansion) != _IDENTITY_PAGE_INDEX_EXPANSION: raise RuntimeError("QSA does not support expanded V2 page indices") layer_offset_pages = int(converter.layer_offset) - page_upper = self.impl.get_page_index_upper_bound(layer_offset, Role.KEY) + page_upper = self.impl.get_page_index_upper_bound(layer_offset, key_role) num_slots_total = page_upper + layer_offset_pages if num_slots_total % scale != 0: raise RuntimeError( @@ -202,17 +348,23 @@ def get_buffers( ) num_slots = num_slots_total // scale - torch_dtype = binding_to_torch_dtype(self.dtype) head_dim = self.head_dim_per_layer[layer_offset] num_heads = self.num_kv_heads_per_layer[layer_offset] + width = head_dim // values_per_element if kv_layout == "NHD": - page_shape = [self.tokens_per_block, num_heads, head_dim] + page_shape = [self.tokens_per_block, num_heads, width] else: - page_shape = [num_heads, self.tokens_per_block, head_dim] + page_shape = [num_heads, self.tokens_per_block, width] + page_bytes = self.tokens_per_block * num_heads * width * dtype.itemsize + if page_bytes != key_stride: + raise RuntimeError( + f"QSA {key_role} page is {key_stride} B, but {page_shape} {dtype} is " + f"{page_bytes} B: layer={layer_idx}" + ) full_view = convert_to_torch_tensor( TensorWrapper( addr_key, - torch_dtype, + dtype, [num_slots, scale, *page_shape], ) ) @@ -220,9 +372,27 @@ def get_buffers( @torch.compiler.disable def get_qsa_position_buffer(self) -> Optional[torch.Tensor]: - """Return per-token three-axis RoPE/mRoPE position coordinates.""" + """Return per-token three-axis RoPE/mRoPE position coordinates. + + int32 ``[slots, tokens_per_block, 3]``, indexed by lifecycle slot. + """ if self.qsa_position_layer_id is None: return None + if self.qsa_side_planes_size is not None: + position_size = ( + QSA_POSITION_COORDINATE_AXES * _POSITION_ELEMENT_BYTES * self.tokens_per_block + ) + planes = self._side_planes() + start = len(self.qsa_local_sparse_layer_ids) * ( + self.qsa_index_kv_heads + * self.qsa_index_dim + * _INDEX_K_ELEMENT_BYTES + * self.tokens_per_block + ) + plane = planes[:, start : start + position_size] + return plane.view(QSA_POSITION_CACHE_DTYPE).view( + planes.shape[0], self.tokens_per_block, QSA_POSITION_COORDINATE_AXES + ) local_idx = self.layer_offsets[self.qsa_position_layer_id] addr = self.impl.get_mem_pool_base_address( local_idx, @@ -269,6 +439,30 @@ def get_qsa_position_buffer(self) -> Optional[torch.Tensor]: ) return full_view[:, _SHARED_POSITION_ROLE_INDEX] + def _build_pool_mapping_tensors(self): + """Refuse an NVFP4 layout whose K/V and block-scale slots differ. + + The attention op reads a layer's K/V pages and their block scales + through one block-offset table: a page index scaled for the K/V pool + addresses the scale pool too. A layer whose two page-index scales + differ would read and write other slots' scales without any error. + """ + if self.dtype == DataType.NVFP4: + for layer_id in self.pp_layers: + local_idx = self.layer_offsets[layer_id] + if self._is_local_mamba_layer(local_idx): + continue + data_scale = self.impl.get_page_index_scale(local_idx, Role.KEY) + scale_scale = self.impl.get_page_index_scale(local_idx, Role.KEY_BLOCK_SCALE) + if data_scale != scale_scale: + raise RuntimeError( + "NVFP4 K/V and block-scale pools hold a different number of pages " + f"per lifecycle slot (layer {layer_id}: {data_scale} vs {scale_scale}); " + "another buffer was coalesced into one of them, and the attention op " + "would address the block scales of other slots" + ) + return super()._build_pool_mapping_tensors() + def get_qsa_attention_pool_layout(self) -> tuple[int, int]: """Return the shared attention-pool index and V2 page-index scale. @@ -307,4 +501,4 @@ def get_qsa_attention_pool_layout(self) -> tuple[int, int]: return pool_id, scale -__all__ = ["QSAMambaHybridCacheManagerV2"] +__all__ = ["QSA_INDEX_POSITION", "QSA_SIDE_PLANES", "QSAMambaHybridCacheManagerV2"] diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/kernels.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/kernels.py index 74834c64fd48..dec4259363a8 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/kernels.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/kernels.py @@ -2178,6 +2178,7 @@ def _qsa_mtp_reuse_topk_kernel( visible_ptr, out_ptr, BLOCK_TOPK: tl.constexpr, + DISPLACE_OLDEST: tl.constexpr, ): """Rebuild a query row's selection from its request's captured row. @@ -2185,6 +2186,10 @@ def _qsa_mtp_reuse_topk_kernel( groups that completed afterwards are written on top of it. They are always larger than every captured ID, because the capture's Top-K searched only ``[0, captured_visible)``, so the row never gains a duplicate. + + A full row has no room for them: they take the slots of the row's last IDs, + or with ``DISPLACE_OLDEST`` the survivors shift to the front and they take + the slots after them, which for an ascending row displaces its smallest IDs. """ row = tl.program_id(0) request = tl.load(request_ptr + row).to(tl.int32) @@ -2197,8 +2202,16 @@ def _qsa_mtp_reuse_topk_kernel( # Valid captured IDs form a prefix; the rest is -1 padding. filled = tl.sum((values >= 0).to(tl.int32)) fresh = tl.minimum(tl.maximum(current - captured, 0), BLOCK_TOPK) - # Append after the prefix when there is room, otherwise displace its tail. + # Append after the prefix when there is room, otherwise displace its tail, or its head + # (the survivors shifted to the front) with DISPLACE_OLDEST. start = tl.minimum(filled, BLOCK_TOPK - fresh) + if DISPLACE_OLDEST: + source = slots + (filled - start) + values = tl.load( + shared_ptr + request.to(tl.int64) * BLOCK_TOPK + source, + mask=source < BLOCK_TOPK, + other=-1, + ) offset = slots - start is_fresh = (offset >= 0) & (offset < fresh) values = tl.where(is_fresh, (captured + offset).to(values.dtype), values) @@ -2212,8 +2225,13 @@ def triton_qsa_mtp_reuse_topk( request_indices: torch.Tensor, visible_blocks: torch.Tensor, output: torch.Tensor, + displace_oldest: bool = False, ) -> torch.Tensor: - """Expand captured per-request selections back to per-query-row selections.""" + """Expand captured per-request selections back to per-query-row selections. + + ``displace_oldest`` is for captured rows in ascending order: a full row gives up + its first IDs, rather than its last, to the groups completed since the capture. + """ rows, block_topk = output.shape if shared_indices.shape[1] != block_topk: raise ValueError( @@ -2239,6 +2257,7 @@ def triton_qsa_mtp_reuse_topk( visible_blocks, output, BLOCK_TOPK=block_topk, + DISPLACE_OLDEST=displace_oldest, num_warps=4, ) return output diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py index fe1a441e3584..4ccfb3ce2928 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py @@ -53,6 +53,7 @@ def __post_init__(self) -> None: self.qsa_shared_topk_indices: Optional[torch.Tensor] = None self.qsa_shared_topk_visible_blocks: Optional[torch.Tensor] = None self.qsa_shared_topk_captured_slots: set[int] = set() + self.qsa_draft_block_table: Optional[torch.Tensor] = None if not isinstance(self.kv_cache_manager, QSAMambaHybridCacheManagerV2): raise TypeError("QSA sparse attention requires QSAMambaHybridCacheManagerV2") if not isinstance(self.sparse_metadata_params, QSASparseMetadataParams): @@ -80,6 +81,7 @@ def __post_init__(self) -> None: self._qsa_attention_pool_id, self._qsa_page_index_scale = ( self.kv_cache_manager.get_qsa_attention_pool_layout() ) + self._init_qsa_draft_block_table() # Packed-row mapping shared by compression, selection, and sparse GQA. self.qsa_req_idx_per_token = self.get_empty( @@ -164,7 +166,11 @@ def __post_init__(self) -> None: # compressed groups were complete when the row was captured, so the # reuse can tell which groups formed afterwards. if self.sparse_metadata_params.mtp_index_share: + # Slots are numbered by the bound manager; the draft loop binds the + # separate draft manager. num_slots = len(self.kv_cache_manager.qsa_shared_topk_slots) + if self.qsa_draft_block_table is not None: + num_slots = max(num_slots, len(self.draft_kv_cache_manager.qsa_shared_topk_slots)) self.qsa_shared_topk_indices = self.get_empty( buffers, ( @@ -184,6 +190,32 @@ def __post_init__(self) -> None: capture_graph=capture_graph, ) + def _init_qsa_draft_block_table(self) -> None: + """Allocate the page table of a separate one-model draft KV cache. + + The draft layers' pages live in the draft manager's pools, which the + target table does not address. This table is its own graph-stable + buffer, so a captured graph reads the target table in the target + forward and this one in the draft loop; ``prepare()`` refreshes both + before every replay. + """ + draft_manager = self.draft_kv_cache_manager + if ( + not isinstance(draft_manager, QSAMambaHybridCacheManagerV2) + or draft_manager.qsa_position_layer_id is None + ): + return + self.qsa_draft_block_table = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_sequences, draft_manager.max_blocks_per_seq), + cache_name="qsa_draft_block_table", + dtype=torch.int32, + capture_graph=self.is_cuda_graph, + ) + self._qsa_draft_attention_pool_id, self._qsa_draft_page_index_scale = ( + draft_manager.get_qsa_attention_pool_layout() + ) + def set_skip_topk(self, skip: bool) -> None: """Mark a draft step that reuses the captured selection.""" self.qsa_indexer_skip_topk = skip @@ -245,21 +277,45 @@ def _refresh_qsa_token_mapping(self) -> None: ((logical + 1) // self.sparse_metadata_params.compress_ratio).to(torch.int32) ) - def _refresh_qsa_block_table(self) -> None: - """Decode V2 attention offsets into raw page IDs used by QSA kernels.""" + def _unscale_qsa_block_table( + self, + block_offsets: torch.Tensor, + pool_id: int, + page_index_scale: int, + block_table: torch.Tensor, + ) -> None: + """Decode one manager's V2 attention offsets into raw page IDs.""" num_seqs = self.num_seqs if num_seqs <= 0: return - scaled = self.kv_cache_block_offsets[ - self._qsa_attention_pool_id, - :num_seqs, - QSA_KEY_ROLE_INDEX, - :, - ] triton_qsa_unscale_block_table( - scaled_block_table=scaled, - block_table=self.qsa_block_table[:num_seqs], - page_index_scale=self._qsa_page_index_scale, + scaled_block_table=block_offsets[pool_id, :num_seqs, QSA_KEY_ROLE_INDEX, :], + block_table=block_table[:num_seqs], + page_index_scale=page_index_scale, + ) + + def _refresh_qsa_block_table(self) -> None: + """Decode the bound manager's offsets into the table QSA kernels read. + + While ``prepare_for_draft_forward()`` is in effect every operand here is + the draft manager's, so this rebuilds the draft table. + """ + self._unscale_qsa_block_table( + self.kv_cache_block_offsets, + self._qsa_attention_pool_id, + self._qsa_page_index_scale, + self.qsa_block_table, + ) + + def _refresh_qsa_draft_block_table(self) -> None: + """Decode the separate draft manager's offsets into the draft table.""" + if self.qsa_draft_block_table is None: + return + self._unscale_qsa_block_table( + self.draft_kv_cache_block_offsets, + self._qsa_draft_attention_pool_id, + self._qsa_draft_page_index_scale, + self.qsa_draft_block_table, ) def _refresh_qsa_speculative_snapshot_state(self) -> None: @@ -276,8 +332,29 @@ def _refresh_qsa_device_state(self) -> None: if not self.qsa_has_local_layers: return self._refresh_qsa_block_table() + self._refresh_qsa_draft_block_table() self._refresh_qsa_token_mapping() + def prepare_for_draft_forward(self) -> None: + """Point the QSA page table at the separate draft manager's pages. + + The caller has already swapped ``kv_cache_manager`` and + ``kv_cache_block_offsets`` to the draft manager's. The draft table was + decoded in ``prepare()``, so this only rebinds fields and launches + nothing; ``restore_draft_swaps()`` restores the target bindings. The + row mapping is manager-independent and stays as it is. + """ + if self.qsa_draft_block_table is None: + return + if self.kv_cache_manager is not self.draft_kv_cache_manager: + raise RuntimeError( + "QSA draft forward expects the metadata's draft KV cache manager " + f"to be bound; got {type(self.kv_cache_manager).__name__}" + ) + self.swap_for_draft("qsa_block_table", self.qsa_draft_block_table) + self.swap_for_draft("_qsa_attention_pool_id", self._qsa_draft_attention_pool_id) + self.swap_for_draft("_qsa_page_index_scale", self._qsa_draft_page_index_scale) + def prepare(self) -> None: """Refresh host decisions, then rebuild graph-stable device mappings.""" super().prepare() diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 84edf64cfd06..9983032e9b68 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -284,6 +284,8 @@ class ModelConfig(Generic[TConfig]): # disables the optimization; hardware support is resolved by the backend. skip_correction_threshold: float = 0.0 moe_backend: str = 'CUTLASS' # options can be CUTLASS, TRTLLM + # The backend the user requested, before AUTO resolution. + moe_backend_requested: str = 'AUTO' # IF true, disables FC2+finalize fusion in CUTLASS MoE backend moe_disable_finalize_fusion: bool = False # If true, use low precision combine in MoE operations (only for NVFP4 quantization) @@ -1526,6 +1528,7 @@ def _recursive_update_config(config: transformers.PretrainedConfig, architecture, quant_config=quant_config, ) + kwargs['moe_backend_requested'] = requested_moe_backend # V4.1 names the routed experts exactly as V4 does and encodes MXFP4 the # same way (I8 storage), so the detector and setter apply verbatim. diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py index ac57b3ef45eb..c204b6d2a508 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py @@ -292,6 +292,9 @@ class HostInputs: n_gen: int = 0 has_any_multimodal_request: bool = False + # Whether a verify group or generation request of this pass carries real + # MRoPE metadata; without one, the delta read slots are dropped. + has_gen_mrope_delta: bool = False beam_width: int = 1 num_tokens: int = 0 @@ -1203,6 +1206,19 @@ def _classify_generation_requests( else: generation_requests.append(request) extend_requests += extend_dummy_requests + # The engine no longer builds first-draft requests; one that still carries the + # flag is refused on an mRoPE model rather than packed without a delta slot. + first_draft_requests = [ + request + for request in scheduled_requests.generation_requests + if getattr(request, "py_is_first_draft", False) + ] + if host.use_mrope and first_draft_requests: + raise NotImplementedError( + "First-draft requests on an mRoPE model are not supported: they add no " + "mRoPE delta read slot, so every later generation request would read " + "another request's delta." + ) def _prepare_helix_state( self, host: HostInputs, kv_cache_manager: KVCacheManager | KVCacheManagerV2 | None @@ -1242,6 +1258,30 @@ def _helix_pack_extend(request, group: int) -> int: host.helix_pack_extend = _helix_pack_extend host.has_cp_helix = _has_cp_helix + @staticmethod + def _decode_mrope_position_delta(request: LlmRequest) -> torch.Tensor | None: + """The MRoPE position delta of a decode-phase request, or None. + + A request that ran its context phase on this worker already holds the + delta in ``py_mrope_position_delta``. Otherwise (for example the first + step of a generation-only request) it is read from the request's + ``mrope_config``, moved to the device once, and cached on the request. + None means the request carries no MRoPE metadata (a text-only prompt), + whose delta is zero. + """ + mrope_position_delta = getattr(request, "py_mrope_position_delta", None) + if mrope_position_delta is None and request.py_multimodal_data: + mrope_config = request.py_multimodal_data.get("mrope_config") or {} + mrope_position_delta = mrope_config.get("mrope_position_deltas") + if mrope_position_delta is not None: + if mrope_position_delta.device.type == "cpu": + mrope_position_delta = maybe_pin_memory(mrope_position_delta).to( + device="cuda", dtype=torch.int32, non_blocking=True + ) + mrope_config["mrope_position_deltas"] = mrope_position_delta + request.py_mrope_position_delta = mrope_position_delta + return mrope_position_delta + def _pack_extend_requests( self, host: HostInputs, @@ -1273,8 +1313,24 @@ def _pack_extend_requests( _has_cp_helix = host.has_cp_helix _helix_local_len_host = host.helix_local_len_host _helix_pack_extend = host.helix_pack_extend + mrope_position_ids = host.lists.mrope_position_ids + mrope_delta_write_seq_slots = host.lists.mrope_delta_write_seq_slots + mrope_delta_read_seq_slots = host.lists.mrope_delta_read_seq_slots + multimodal_params_list = host.lists.multimodal_params_list + _use_mrope = host.use_mrope + mrope_dummy_seq_slot = host.mrope_dummy_seq_slot + has_gen_mrope_delta = host.has_gen_mrope_delta runtime_draft_token_buffer_width = runtime_tokens_per_gen_step - 1 + # Under extend_ctx the attention kernel runs extend requests as chunked + # context rows: their rotary comes from position_ids alone, and they take + # no entry in the generation-indexed MRoPE delta read slots. + extend_requests_are_context = ( + _use_mrope + and enable_spec_decode + and spec_config.spec_dec_mode.extend_ctx(self._config.attention_backend) + and spec_config.is_linear_tree + ) for request in extend_requests: is_promoted_context = request.py_request_id in promoted_context_request_ids if getattr(request, "py_needs_onehot_draft_probs", False): @@ -1406,8 +1462,61 @@ def _pack_extend_requests( else: prompt_lengths.append(request.py_prompt_len) + if _use_mrope: + mrope_position_delta = self._decode_mrope_position_delta(request) + if mrope_position_delta is not None: + # Every row of a verify group lies past the prompt, so it is + # text: its three MRoPE coordinates are equal, the scalar + # position plus the request's delta. The overlap scheduler's + # accepted-token correction later adds to all three axes. + num_request_tokens = sequence_lengths[-1] + token_start = len(position_ids) - num_request_tokens + first_position = position_ids[token_start] + extend_mrope_position_ids = ( + torch.arange( + first_position, + first_position + num_request_tokens, + dtype=torch.int32, + device=mrope_position_delta.device, + ) + + mrope_position_delta.reshape(1, 1, 1) + ).expand(3, 1, num_request_tokens) + mrope_position_ids.append( + (token_start, len(position_ids), extend_mrope_position_ids) + ) + if ( + request.py_seq_slot is not None + and not request.is_dummy + and getattr(request, "py_mrope_delta_cache_slot", None) + != request.py_seq_slot + ): + multimodal_params_list.append( + MultimodalParams( + multimodal_data={ + "mrope_config": {"mrope_position_deltas": mrope_position_delta} + } + ) + ) + mrope_delta_write_seq_slots.append(request.py_seq_slot) + request.py_mrope_delta_cache_slot = request.py_seq_slot + if not extend_requests_are_context: + # The kernel indexes deltas by generation request, so every + # verify group takes one read slot, in batch order; a dummy, + # unslotted, or text-only request reads the zero slot. + if mrope_position_delta is not None: + has_gen_mrope_delta = True + mrope_delta_read_seq_slots.append( + request.py_seq_slot + if mrope_position_delta is not None + and not request.is_dummy + and request.py_seq_slot is not None + else mrope_dummy_seq_slot + ) + extra_inputs.add_generation_request(request) + host.has_gen_mrope_delta = has_gen_mrope_delta + def _pack_generation_requests( self, host: HostInputs, @@ -1439,9 +1548,9 @@ def _pack_generation_requests( _use_mrope = host.use_mrope _has_cp_helix = host.has_cp_helix mrope_dummy_seq_slot = host.mrope_dummy_seq_slot - # Whether any generation request in this batch carries real MRoPE - # metadata; see the post-loop cleanup below. - has_gen_mrope_delta = False + # Whether any generation request in this batch, verify groups included, + # carries real MRoPE metadata; see the post-loop cleanup below. + has_gen_mrope_delta = host.has_gen_mrope_delta # Cache invariant method result to avoid repeated calls per-request _n_gen = len(generation_requests) @@ -1572,17 +1681,7 @@ def _pack_generation_requests( gather_ids.append(len(position_ids) - 1) if _use_mrope: - mrope_position_delta = getattr(request, "py_mrope_position_delta", None) - if mrope_position_delta is None and request.py_multimodal_data: - mrope_config = request.py_multimodal_data.get("mrope_config") or {} - mrope_position_delta = mrope_config.get("mrope_position_deltas") - if mrope_position_delta is not None: - if mrope_position_delta.device.type == "cpu": - mrope_position_delta = maybe_pin_memory(mrope_position_delta).to( - device="cuda", dtype=torch.int32, non_blocking=True - ) - mrope_config["mrope_position_deltas"] = mrope_position_delta - request.py_mrope_position_delta = mrope_position_delta + mrope_position_delta = self._decode_mrope_position_delta(request) if mrope_position_delta is not None: has_gen_mrope_delta = True # NOTE: Expanding position_ids to 3D tensor who is using mrope diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 9f2f4e12a855..6cecef79e926 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -133,16 +133,22 @@ def _validate_and_adjust_mamba_snapshot_config(config: ModelConfig, kv_cache_config.enable_block_reuse = False +def effective_mamba_ssm_cache_dtype(config: ModelConfig, + mamba_ssm_cache_dtype: str) -> torch.dtype: + """The recurrent-state cache dtype the engine runs with.""" + if mamba_ssm_cache_dtype == "auto": + return resolve_auto_ssm_cache_dtype(config.pretrained_config, + config.torch_dtype) + return str_dtype_to_torch(mamba_ssm_cache_dtype) + + def validate_and_set_mamba_ssm_cache_dtype( config: ModelConfig, mamba_ssm_cache_dtype: str, mamba_ssm_stochastic_rounding: bool = False, mamba_ssm_philox_rounds: int = 10) -> None: - if mamba_ssm_cache_dtype == "auto": - mamba_ssm_cache_dtype = resolve_auto_ssm_cache_dtype( - config.pretrained_config, config.torch_dtype) - else: - mamba_ssm_cache_dtype = str_dtype_to_torch(mamba_ssm_cache_dtype) + mamba_ssm_cache_dtype = effective_mamba_ssm_cache_dtype( + config, mamba_ssm_cache_dtype) validate_kimi_kda_state_dtype(config.pretrained_config, mamba_ssm_cache_dtype, mamba_ssm_stochastic_rounding) @@ -152,6 +158,15 @@ def validate_and_set_mamba_ssm_cache_dtype( config.quant_config.mamba_ssm_philox_rounds = mamba_ssm_philox_rounds +def effective_kv_cache_quant(model_config: ModelConfig, + pyt_kv_cache_dtype: str): + """The KV cache quantization the engine runs with: the checkpoint's under + "auto", otherwise what the explicit dtype maps to.""" + if pyt_kv_cache_dtype == "auto": + return model_config.quant_config.kv_cache_quant_algo + return _KV_CACHE_MAP.get(pyt_kv_cache_dtype, None) + + def validate_and_set_kv_cache_quant(model_config: ModelConfig, pyt_kv_cache_dtype: str) -> None: logger.info( @@ -163,9 +178,8 @@ def validate_and_set_kv_cache_quant(model_config: ModelConfig, valid_pyt_quant = bool(pyt_kv_cache_dtype in _VALID_KV_CACHE_DTYPES) mapped_pyt_quant = _KV_CACHE_MAP.get(pyt_kv_cache_dtype, None) - effective_kv_cache_quant = (kv_cache_quant if pyt_kv_cache_dtype == "auto" - else mapped_pyt_quant) - if (effective_kv_cache_quant in (QuantAlgo.NVFP4, QuantAlgo.NVFP4.value) + effective_quant = effective_kv_cache_quant(model_config, pyt_kv_cache_dtype) + if (effective_quant in (QuantAlgo.NVFP4, QuantAlgo.NVFP4.value) and not supports_fp4_mla_attention(model_config) and torch.cuda.is_available() and get_sm_version() == 107): logger.warning( diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 9455105c8029..5b361da5bea8 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -6496,6 +6496,12 @@ def _validate_request(self, request: LlmRequest): # Perform sampler-specific validation self.sampler.validate_request(request) + # Perform model-specific validation, for models that define the hook + validate = getattr(getattr(self.model_engine, "model", None), + "validate_request", None) + if validate is not None: + validate(request) + # Check if request has enough budget self._validate_request_budget(request) diff --git a/tests/unittest/_torch/attention/sparse/qsa/test_qsa_cache_manager_layout.py b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_cache_manager_layout.py new file mode 100644 index 000000000000..5c811ea87d5b --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_cache_manager_layout.py @@ -0,0 +1,224 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The QSA hybrid cache manager's V2 page layout, on real managers. + +The attention op addresses a layer's K/V page as ``pool + (scale * slot + 2 * +offset) * page`` and its NVFP4 block-scale page with the same block index in +the scale pool. These tests build the managers a Qwen3.8-Flash-Next engine +builds -- the 48-layer trunk (12 QSA layers among 36 recurrent ones) and the +one-layer MTP draft manager -- and check that every view the manager serves is +where that formula lands, for bf16, fp8 and NVFP4 pools. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.sparse.qsa import QSAMambaHybridCacheManagerV2 +from tensorrt_llm._torch.attention.backends.sparse.qsa.cache_manager import ( + QSA_INDEX_POSITION, + QSA_SIDE_PLANES, +) +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import Role +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp +from tensorrt_llm.llmapi.llm_args import KvCacheConfig, QSASparseAttentionConfig +from tensorrt_llm.mapping import Mapping + +_CHECKPOINT = SimpleNamespace( + indexer_n_heads=4, + indexer_kv_heads=1, + indexer_head_dim=128, + indexer_budget=2048, + indexer_compress_ratio=4, +) +_TRUNK_LAYERS = 48 +_QSA_LAYERS = list(range(3, _TRUNK_LAYERS, 4)) +_DRAFT_LAYER = _TRUNK_LAYERS +_KV_HEADS, _HEAD_DIM, _TPB = 2, 256, 32 +_INDEX_PAGE = 1 * 128 * 2 * _TPB +_POSITION_PAGE = 3 * 4 * _TPB + + +def _manager(dtype: DataType, *, draft: bool) -> QSAMambaHybridCacheManagerV2: + """The trunk or the separate MTP draft manager, as the engine builds them.""" + if draft: + attention_mask = [False] * _TRUNK_LAYERS + [True] + mamba_mask = [False] * (_TRUNK_LAYERS + 1) + else: + attention_mask = [layer in _QSA_LAYERS for layer in range(_TRUNK_LAYERS)] + mamba_mask = [not attention for attention in attention_mask] + return QSAMambaHybridCacheManagerV2( + mamba_d_state=8, + mamba_d_conv=4, + mamba_num_heads=4, + mamba_n_groups=1, + mamba_head_dim=8, + mamba_num_layers=sum(mamba_mask), + mamba_layer_mask=mamba_mask, + mamba_cache_dtype=torch.bfloat16, + mamba_ssm_cache_dtype=torch.float32, + kv_cache_config=KvCacheConfig(max_tokens=4096, enable_block_reuse=False), + kv_cache_type=CacheTypeCpp.SELF, + num_layers=sum(attention_mask), + num_kv_heads=_KV_HEADS, + head_dim=_HEAD_DIM, + tokens_per_block=_TPB, + max_seq_len=1024, + max_batch_size=4, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=dtype, + layer_mask=attention_mask, + is_draft=draft, + vocab_size=1024, + conv_state_layout="q_k_v", + sparse_attention_config=QSASparseAttentionConfig(), + pretrained_config=_CHECKPOINT, + ) + + +def _layers(draft: bool) -> list[int]: + return [_DRAFT_LAYER] if draft else _QSA_LAYERS + + +def _kv_page_bytes(dtype: DataType) -> int: + element_bits = {DataType.BF16: 16, DataType.FP8: 8, DataType.NVFP4: 4}[dtype] + return _KV_HEADS * _TPB * _HEAD_DIM * element_bits // 8 + + +def _thop_addresses(cm, layer: int, slot: int, page_bytes: int, plane: int) -> tuple[int, int]: + """Byte addresses of this layer's K and V page for `slot`, as the attention + op forms them from the manager's pool pointers, mapping and page-index + scale (`plane` 0 = data, 1 = NVFP4 block scales).""" + local = cm.layer_offsets[layer] + pool, offset = cm.kv_cache_pool_mapping[local].tolist() + pointers = cm.kv_cache_pool_pointers + base = int(pointers[pool, 0, plane]) if pointers.dim() == 3 else int(pointers[pool, 0]) + key_entry = int(cm.index_scales[pool]) * slot + value_entry = key_entry + int(cm.kv_offset[pool]) + shift = offset * cm.kv_factor * page_bytes + return base + key_entry * page_bytes + shift, base + value_entry * page_bytes + shift + + +def _slots_to_check(view: torch.Tensor) -> list[int]: + return [0, 1, view.shape[0] // 2, view.shape[0] - 1] + + +@pytest.fixture(scope="module", autouse=True) +def _require_cuda(): + assert torch.cuda.is_available(), "the managers allocate their pools on the GPU" + + +@pytest.mark.parametrize("draft", [False, True], ids=["trunk", "draft"]) +@pytest.mark.parametrize("dtype", [DataType.BF16, DataType.FP8, DataType.NVFP4]) +def test_views_are_where_the_attention_op_addresses(dtype: DataType, draft: bool) -> None: + """Every K/V (and NVFP4 block-scale) view of every QSA layer starts, at + every checked slot, at the address the attention op computes for it.""" + cm = _manager(dtype, draft=draft) + try: + page = _kv_page_bytes(dtype) + for layer in _layers(draft): + kv = cm.get_buffers(layer, kv_layout="HND") + assert kv.shape[1:] == ( + 2, + _KV_HEADS, + _TPB, + _HEAD_DIM // (2 if dtype == DataType.NVFP4 else 1), + ) + for slot in _slots_to_check(kv): + key, value = _thop_addresses(cm, layer, slot, page, plane=0) + assert kv[slot, 0].data_ptr() == key, (layer, slot) + assert kv[slot, 1].data_ptr() == value, (layer, slot) + if dtype != DataType.NVFP4: + continue + scales = cm.get_block_scale_buffers(layer, kv_layout="HND") + assert scales.dtype == torch.uint8 + assert scales.shape == (kv.shape[0], 2, _KV_HEADS, _TPB, _HEAD_DIM // 16) + for slot in _slots_to_check(scales): + key, value = _thop_addresses(cm, layer, slot, page // 8, plane=1) + assert scales[slot, 0].data_ptr() == key, (layer, slot) + assert scales[slot, 1].data_ptr() == value, (layer, slot) + finally: + cm.shutdown() + + +@pytest.mark.parametrize("draft", [False, True], ids=["trunk", "draft"]) +def test_nvfp4_scale_pages_follow_the_data_pages(draft: bool) -> None: + """Under NVFP4 the K/V pool and the block-scale pool hold the same number + of pages per slot, and the index and position planes are packed into a + buffer of their own instead of landing in the K/V pool.""" + cm = _manager(DataType.NVFP4, draft=draft) + try: + layers = _layers(draft) + assert cm.qsa_side_planes_size == len(layers) * _INDEX_PAGE + _POSITION_PAGE + for ordinal, layer in enumerate(layers): + local = cm.layer_offsets[layer] + data_scale = cm.impl.get_page_index_scale(local, Role.KEY) + assert data_scale == cm.impl.get_page_index_scale(local, Role.KEY_BLOCK_SCALE) + assert data_scale == 2 * len(layers) + assert cm.kv_cache_pool_mapping[local, 1] == ordinal + planes_owner = cm.layer_offsets[layers[0]] + assert cm.impl.get_page_index_scale(planes_owner, QSA_SIDE_PLANES) == 1 + finally: + cm.shutdown() + + +@pytest.mark.parametrize("draft", [False, True], ids=["trunk", "draft"]) +@pytest.mark.parametrize("dtype", [DataType.BF16, DataType.FP8]) +def test_bf16_and_fp8_keep_per_layer_side_roles(dtype: DataType, draft: bool) -> None: + """bf16 and fp8 pages never share a size with an index page, so the side + planes stay the per-layer INDEX_KEY role plus one position role, served by + the same views the generic accessor builds.""" + cm = _manager(dtype, draft=draft) + try: + layers = _layers(draft) + assert cm.qsa_side_planes_size is None + for layer in layers: + local = cm.layer_offsets[layer] + assert cm.impl.get_page_index_scale(local, Role.INDEX_KEY) == len(layers) + assert cm.impl.get_page_index_scale(local, Role.KEY) == 2 * len(layers) + generic = MambaHybridCacheManagerV2.get_index_k_buffer( + cm, layer, num_heads=1, head_dim=128, dtype=torch.bfloat16 + ) + served = cm.get_index_k_buffer(layer) + assert (served.data_ptr(), served.shape, served.stride()) == ( + generic.data_ptr(), + generic.shape, + generic.stride(), + ) + owner = cm.layer_offsets[layers[0]] + assert cm.impl.get_page_index_scale(owner, QSA_INDEX_POSITION) == 1 + assert cm.get_qsa_position_buffer().stride(0) == _TPB * 3 + finally: + cm.shutdown() + + +@pytest.mark.parametrize("draft", [False, True], ids=["trunk", "draft"]) +@pytest.mark.parametrize("dtype", [DataType.BF16, DataType.FP8, DataType.NVFP4]) +def test_side_and_main_views_do_not_alias(dtype: DataType, draft: bool) -> None: + """One slot written through every view reads back unchanged through each: + no K/V, scale, index or position page overlaps another.""" + cm = _manager(dtype, draft=draft) + try: + views = [] + for layer in _layers(draft): + kv = cm.get_buffers(layer, kv_layout="HND") + views += [kv[:, 0], kv[:, 1]] + if dtype == DataType.NVFP4: + scales = cm.get_block_scale_buffers(layer, kv_layout="HND") + views += [scales[:, 0], scales[:, 1]] + views.append(cm.get_index_k_buffer(layer)) + views.append(cm.get_qsa_position_buffer()) + slots = {view.shape[0] for view in views} + assert len(slots) == 1, f"views disagree on the slot count: {sorted(slots)}" + slot = slots.pop() - 1 + as_bytes = [view[slot].view(torch.uint8) for view in views] + for i, raw in enumerate(as_bytes): + raw.fill_(i + 1) + torch.cuda.synchronize() + for i, raw in enumerate(as_bytes): + assert bool((raw == i + 1).all()), f"view {i} was overwritten by another view" + finally: + cm.shutdown() diff --git a/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py index b766051d0ef0..6dc868d1c5fb 100644 --- a/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py +++ b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py @@ -11,6 +11,7 @@ QSASparseMetadataParams, QSASparseParams, ) +from tensorrt_llm._torch.attention.backends.sparse.qsa.cache_manager import QSA_SIDE_PLANES from tensorrt_llm._torch.attention.backends.sparse.qsa.indexer import ( QSAIndexer, _position_coordinates, @@ -37,6 +38,11 @@ select_qsa_paged_tokens, select_qsa_tokens, ) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.speculative.interface import ( + prepare_attn_metadata_for_draft_replay, + restore_attn_metadata_after_draft_replay, +) from tensorrt_llm.runtime.kv_cache_manager_v2 import PageIndexMode @@ -219,6 +225,33 @@ def test_qsa_mtp_reuse_topk_appends_groups_completed_since_capture() -> None: assert row.tolist()[: len(valid)] == valid +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_qsa_mtp_reuse_topk_displace_oldest_shifts_a_full_row() -> None: + """With `displace_oldest` a full row gives up its first ids and the rest move to the + front, so an ascending row stays ascending and keeps its newest groups.""" + shared = torch.tensor( + [[1, 2, 4, 6, 7, 9, 11, 12], [1, 2, 4, 6, 7, 9, -1, -1], [3, 5, 6, 8, 9, 10, 11, -1]], + device="cuda", + dtype=torch.int32, + ) + out = torch.full((3, 8), -99, device="cuda", dtype=torch.int32) + + triton_qsa_mtp_reuse_topk( + shared_indices=shared, + shared_visible_blocks=torch.tensor([13, 10, 12], device="cuda", dtype=torch.int32), + request_indices=torch.tensor([0, 1, 2], device="cuda", dtype=torch.int32), + visible_blocks=torch.tensor([15, 11, 14], device="cuda", dtype=torch.int32), + output=out, + displace_oldest=True, + ) + + assert out.tolist() == [ + [4, 6, 7, 9, 11, 12, 13, 14], # full, 2 new: the 2 oldest go + [1, 2, 4, 6, 7, 9, 10, -1], # room: appended + [5, 6, 8, 9, 10, 11, 12, 13], # 7 ids, 2 new: the oldest goes + ] + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") def test_qsa_mtp_reuse_topk_is_verbatim_without_new_groups() -> None: """A draft step that completes no group reuses the captured row unchanged.""" @@ -1456,12 +1489,39 @@ def test_qsa_side_buffers_use_exact_geometry_and_one_position_role() -> None: manager.qsa_index_kv_heads = 1 manager.qsa_sparse_layer_ids = [7, 11] manager.layer_offsets = {7: 2, 11: 4} + # bf16 K/V pages of 2 heads x 128 dims: no side page shares their size. + manager._main_cache_page_sizes = lambda tokens_per_block: {2 * 128 * 2 * tokens_per_block} buffers = manager._extra_buffers_per_layer(tokens_per_block=128) assert buffers[2][0].size == 128 * 2 * 128 assert len(buffers[2]) == 2 assert len(buffers[4]) == 1 + assert manager.qsa_side_planes_size is None + + +def test_qsa_side_buffers_pack_when_a_page_size_matches_the_main_cache() -> None: + """An index page the size of a K/V page would coalesce into the K/V pool, + so both sparse layers' index pages and the position page become one + buffer on the first layer, sized past every main-cache page size.""" + manager = object.__new__(QSAMambaHybridCacheManagerV2) + manager.qsa_index_dim = 128 + manager.qsa_index_kv_heads = 1 + manager.qsa_sparse_layer_ids = [7, 11] + manager.layer_offsets = {7: 2, 11: 4} + index_page, position_page = 128 * 2 * 32, 3 * 4 * 32 + packed = 2 * index_page + position_page + # NVFP4 at 32 tokens per block: K == V == index page; the scale page, and + # a page that happens to equal the packed size, force it to grow once. + manager._main_cache_page_sizes = lambda tokens_per_block: {index_page, 1024, packed} + + buffers = manager._extra_buffers_per_layer(tokens_per_block=32) + + assert list(buffers) == [2] + (packed_buffer,) = buffers[2] + assert packed_buffer.role == QSA_SIDE_PLANES + assert packed_buffer.size == packed + 128 + assert manager.qsa_side_planes_size == packed + 128 def test_qsa_position_buffer_keeps_shared_index_mode_outside_dynamo() -> None: @@ -1478,6 +1538,7 @@ def get_page_stride(self, layer_idx, role): manager = object.__new__(QSAMambaHybridCacheManagerV2) manager.qsa_position_layer_id = 7 + manager.qsa_side_planes_size = None manager.layer_offsets = {7: 2} manager.impl = _Impl() manager.tokens_per_block = 128 @@ -1608,3 +1669,241 @@ def test_qsa_multi_token_generation_captures_speculative_cache_snapshot() -> Non assert pending["token_ordinals"].tolist() == [0, 1, 2, 3] torch.testing.assert_close(pending["index_values"], index_cache[pages, within]) torch.testing.assert_close(pending["position_values"], position_cache[pages, within]) + + +class _QsaPageConverter: + def __init__(self, scale: int) -> None: + self.scale = scale + self.expansion = 1 + + +class _QsaPageImpl: + def __init__(self, scale: int) -> None: + self._scale = scale + + def get_page_index_converter(self, layer_idx, role): + del layer_idx, role + return _QsaPageConverter(self._scale) + + +def _qsa_layout_manager( + *, sparse_layer_ids, pool_id, scale, num_pools, max_blocks_per_seq +) -> QSAMambaHybridCacheManagerV2: + """A V2 QSA manager reduced to the page layout the metadata reads.""" + manager = object.__new__(QSAMambaHybridCacheManagerV2) + manager.qsa_sparse_layer_ids = list(sparse_layer_ids) + manager.layer_offsets = {layer: local for local, layer in enumerate(sparse_layer_ids)} + manager.qsa_position_layer_id = sparse_layer_ids[0] + manager.layer_to_pool_mapping_dict = {local: pool_id for local in range(len(sparse_layer_ids))} + manager.num_attention_op_pools = num_pools + manager.max_blocks_per_seq = max_blocks_per_seq + manager.impl = _QsaPageImpl(scale) + manager.host_kv_cache_block_offsets = torch.zeros( + (num_pools, 4, 2, max_blocks_per_seq), dtype=torch.int32 + ) + return manager + + +# Target: two sparse layers in pool 1 at scale 6. Draft: one MTP layer in pool 0 at +# scale 4, a narrower table, and slot numbers the target never uses. +_TARGET_LAYOUT = dict( + sparse_layer_ids=[3, 7], pool_id=1, scale=6, num_pools=2, max_blocks_per_seq=8 +) +_DRAFT_LAYOUT = dict(sparse_layer_ids=[48], pool_id=0, scale=4, num_pools=1, max_blocks_per_seq=4) + + +def _qsa_metadata_with_draft( + monkeypatch, + *, + mtp_index_share: bool = False, + draft_layout: dict | None = _DRAFT_LAYOUT, + draft_manager=None, +): + """QSA metadata over a target manager and, unless `draft_layout` is None, a separate + draft manager (`draft_manager` replaces the one `draft_layout` describes).""" + monkeypatch.setattr(TrtllmAttentionMetadata, "__post_init__", lambda self: None) + metadata = object.__new__(QSAAttentionMetadata) + metadata.kv_cache_manager = _qsa_layout_manager(**_TARGET_LAYOUT) + if draft_manager is None and draft_layout is not None: + draft_manager = _qsa_layout_manager(**draft_layout) + metadata.draft_kv_cache_manager = draft_manager + metadata.sparse_metadata_params = QSASparseMetadataParams( + token_topk=8, compress_ratio=4, mtp_index_share=mtp_index_share + ) + metadata.max_num_sequences = 4 + metadata.max_num_tokens = 16 + metadata.draft_replay_swapped_attrs = {} + metadata.__post_init__() + return metadata + + +def _set_qsa_offsets(metadata, num_seqs: int, generation: int): + """Write fresh target and draft attention offsets; return the raw slots each + manager's table must decode to.""" + slots = [] + for layout, base in ((_TARGET_LAYOUT, 0), (_DRAFT_LAYOUT, 1000)): + width = layout["max_blocks_per_seq"] + table = torch.arange(num_seqs * width, dtype=torch.int32, device="cuda") + table = table.reshape(num_seqs, width) + base + 100 * generation + table[:, -1] = -1 + # V2 offsets: the key role of the layer's pool holds slot * scale, unallocated + # entries stay -1, and every other pool and role is noise. + offsets = torch.full( + (layout["num_pools"], num_seqs, 2, width), 7777, dtype=torch.int32, device="cuda" + ) + offsets[layout["pool_id"], :, 0, :] = torch.where(table < 0, table, table * layout["scale"]) + slots.append(table) + if layout is _TARGET_LAYOUT: + metadata.kv_cache_block_offsets = offsets + else: + metadata.draft_kv_cache_block_offsets = offsets + return slots + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_qsa_draft_block_table_exists_only_for_a_separate_draft_manager(monkeypatch) -> None: + metadata = _qsa_metadata_with_draft(monkeypatch) + + # Each table is sized, and its pool and page scale taken, from its own manager. + assert metadata.qsa_block_table.shape == (4, _TARGET_LAYOUT["max_blocks_per_seq"]) + assert metadata.qsa_draft_block_table.shape == (4, _DRAFT_LAYOUT["max_blocks_per_seq"]) + assert metadata.qsa_draft_block_table.data_ptr() != metadata.qsa_block_table.data_ptr() + assert (metadata._qsa_attention_pool_id, metadata._qsa_page_index_scale) == (1, 6) + assert (metadata._qsa_draft_attention_pool_id, metadata._qsa_draft_page_index_scale) == (0, 4) + + target_only = _qsa_metadata_with_draft(monkeypatch, draft_layout=None) + assert target_only.qsa_draft_block_table is None + assert target_only.prepare_for_draft_forward() is None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("kind", ("not_qsa", "no_local_layer")) +def test_qsa_draft_table_needs_a_qsa_draft_manager_with_local_layers(monkeypatch, kind) -> None: + if kind == "not_qsa": + draft_manager = SimpleNamespace(qsa_position_layer_id=0) + else: + draft_manager = _qsa_layout_manager(**_DRAFT_LAYOUT) + draft_manager.qsa_position_layer_id = None + metadata = _qsa_metadata_with_draft(monkeypatch, draft_manager=draft_manager) + metadata._seq_lens = torch.ones(2, dtype=torch.int32) + + assert metadata.qsa_draft_block_table is None + # The target's own manager is bound, as outside a draft loop, and the hook and the + # refresh both leave the target bindings alone. + metadata.prepare_for_draft_forward() + metadata._refresh_qsa_draft_block_table() + assert metadata.draft_replay_swapped_attrs == {} + assert (metadata._qsa_attention_pool_id, metadata._qsa_page_index_scale) == (1, 6) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize( + "draft_layout, slots", + ( + (dict(_DRAFT_LAYOUT, sparse_layer_ids=[48, 49, 50]), 3), # draft wider than target + (_DRAFT_LAYOUT, 2), # target wider than draft + (None, 2), # no separate draft manager + ), + ids=("draft_wider", "target_wider", "no_draft"), +) +def test_qsa_shared_topk_slots_cover_the_wider_manager(monkeypatch, draft_layout, slots) -> None: + """The shared selection buffers are addressed by whichever manager is bound.""" + metadata = _qsa_metadata_with_draft( + monkeypatch, mtp_index_share=True, draft_layout=draft_layout + ) + + assert metadata.qsa_shared_topk_indices.shape[0] == slots + assert metadata.qsa_shared_topk_visible_blocks.shape[0] == slots + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_qsa_draft_forward_addresses_the_draft_manager_pages(monkeypatch) -> None: + """prepare() decodes both tables; the draft swap makes `qsa_block_table` the draft one + and the restore puts the target's table, pool and scale back.""" + num_seqs = 3 + metadata = _qsa_metadata_with_draft(monkeypatch) + metadata._seq_lens = torch.ones(num_seqs, dtype=torch.int32) + target_slots, draft_slots = _set_qsa_offsets(metadata, num_seqs, generation=0) + target_table = metadata.qsa_block_table + + metadata._refresh_qsa_block_table() + metadata._refresh_qsa_draft_block_table() + torch.testing.assert_close(target_table[:num_seqs], target_slots) + torch.testing.assert_close(metadata.qsa_draft_block_table[:num_seqs], draft_slots) + + draft_manager = metadata.draft_kv_cache_manager + saved = prepare_attn_metadata_for_draft_replay(metadata, draft_manager) + assert saved is not None + assert metadata.kv_cache_manager is draft_manager + assert metadata.qsa_block_table is metadata.qsa_draft_block_table + # Every operand of the table refresh is the draft manager's while swapped, so a + # refresh rebuilds the draft table and leaves the target's alone. + metadata._refresh_qsa_block_table() + torch.testing.assert_close(metadata.qsa_block_table[:num_seqs], draft_slots) + torch.testing.assert_close(target_table[:num_seqs], target_slots) + + restore_attn_metadata_after_draft_replay(metadata, saved) + assert metadata.kv_cache_manager is not draft_manager + assert metadata.qsa_block_table is target_table + assert (metadata._qsa_attention_pool_id, metadata._qsa_page_index_scale) == (1, 6) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_qsa_prepare_decodes_the_target_and_the_draft_table(monkeypatch) -> None: + """`prepare()` is what refreshes both tables before every step and replay: a second + generation of offsets reaches the target table and the draft table through it.""" + num_seqs = 3 + metadata = _qsa_metadata_with_draft(monkeypatch) + metadata._seq_lens = torch.ones(num_seqs, dtype=torch.int32) + # The parent's prepare and the row mapping are not under test here. + monkeypatch.setattr(TrtllmAttentionMetadata, "prepare", lambda self: None) + monkeypatch.setattr(QSAAttentionMetadata, "_refresh_qsa_token_mapping", lambda self: None) + monkeypatch.setattr( + QSAAttentionMetadata, "_refresh_qsa_speculative_snapshot_state", lambda self: None + ) + for generation in (0, 1): + target_slots, draft_slots = _set_qsa_offsets(metadata, num_seqs, generation) + metadata.prepare() + torch.testing.assert_close(metadata.qsa_block_table[:num_seqs], target_slots) + torch.testing.assert_close(metadata.qsa_draft_block_table[:num_seqs], draft_slots) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_qsa_draft_forward_refuses_an_unswapped_manager(monkeypatch) -> None: + metadata = _qsa_metadata_with_draft(monkeypatch) + + with pytest.raises(RuntimeError, match="draft KV cache manager"): + metadata.prepare_for_draft_forward() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_qsa_draft_table_replays_the_refreshed_draft_pages(monkeypatch) -> None: + """A graph captured across the draft swap reads each table by address: the trunk + segment the target table, the draft segment the draft one, and a refresh outside + capture (as prepare() does before every replay) is what the replay sees.""" + num_seqs = 3 + metadata = _qsa_metadata_with_draft(monkeypatch) + metadata._seq_lens = torch.ones(num_seqs, dtype=torch.int32) + _set_qsa_offsets(metadata, num_seqs, generation=0) + metadata._refresh_qsa_block_table() + metadata._refresh_qsa_draft_block_table() + draft_manager = metadata.draft_kv_cache_manager + target_seen = torch.zeros_like(metadata.qsa_block_table[:num_seqs]) + draft_seen = torch.zeros_like(metadata.qsa_draft_block_table[:num_seqs]) + + graph = torch.cuda.CUDAGraph() + torch.cuda.synchronize() + with torch.cuda.graph(graph): + target_seen.copy_(metadata.qsa_block_table[:num_seqs]) + saved = prepare_attn_metadata_for_draft_replay(metadata, draft_manager) + draft_seen.copy_(metadata.qsa_block_table[:num_seqs]) + restore_attn_metadata_after_draft_replay(metadata, saved) + + target_slots, draft_slots = _set_qsa_offsets(metadata, num_seqs, generation=1) + metadata._refresh_qsa_block_table() + metadata._refresh_qsa_draft_block_table() + graph.replay() + torch.cuda.synchronize() + + torch.testing.assert_close(target_seen, target_slots) + torch.testing.assert_close(draft_seen, draft_slots) diff --git a/tests/unittest/_torch/executor/test_mrope_first_draft_refusal.py b/tests/unittest/_torch/executor/test_mrope_first_draft_refusal.py new file mode 100644 index 000000000000..26947a2e5782 --- /dev/null +++ b/tests/unittest/_torch/executor/test_mrope_first_draft_refusal.py @@ -0,0 +1,54 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A first-draft request in a speculative batch of an mRoPE model is refused. + +The first-draft packing adds neither an mRoPE position span nor a delta read slot, and the +attention kernel indexes the deltas by generation batch index, so every generation request +packed after it would read another request's delta. The engine builds such a request nowhere +today; the runner says so loudly instead of misaligning the slots if one ever appears. +""" + +from unittest.mock import Mock + +import pytest +import test_pytorch_model_engine as engine_tests # module import: its TestCase is not collected + +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm.llmapi import MTPDecodingConfig + + +def _prepare(first_draft: bool) -> None: + case = engine_tests.PyTorchModelEngineTestCase() + model_engine, kv_cache_manager, attn_metadata = case._setup_mrope_engine( + spec_config=MTPDecodingConfig(max_draft_len=1) + ) + try: + # A generation request with an mRoPE delta, then a plain one after it. + first = case._make_mrope_gen_request(4, 1, 0, 10) + first.py_is_first_draft = first_draft + plain = case._make_mrope_gen_request(6, 2, 1, 20) + scheduled_requests = ScheduledRequests() + scheduled_requests.context_requests_last_chunk = [] + scheduled_requests.generation_requests = [first, plain] + model_engine._runner._inputs.prepare_inputs( + scheduled_requests=scheduled_requests, + kv_cache_manager=kv_cache_manager, + attn_metadata=attn_metadata, + spec_metadata=Mock( + _force_non_greedy_for_capture=False, context_prompt_lookahead_tokens=None + ), + enable_spec_decode=True, + runtime_draft_len=1, + is_dummy=False, + ) + finally: + kv_cache_manager.shutdown() + + +def test_first_draft_request_on_an_mrope_model_is_refused() -> None: + with pytest.raises(NotImplementedError, match="First-draft requests on an mRoPE model"): + _prepare(first_draft=True) + + +def test_the_same_batch_without_a_first_draft_request_is_prepared() -> None: + _prepare(first_draft=False) diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index aa38d11c9390..0f204af1154c 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -529,6 +529,92 @@ def test_model_without_remote_tail_does_not_inspect_request() -> None: PyExecutor._configure_csa2_remote_tail(executor, types.SimpleNamespace()) +def _make_admission_executor(model, requests): + """An executor whose `_fetch_and_activate_new_requests` and `_validate_request` are the + real ones, over a model engine holding `model`; everything else they touch is a stub.""" + executor = types.SimpleNamespace( + model_engine=types.SimpleNamespace(model=model), + kv_cache_manager=None, + kv_cache_transceiver=None, + max_beam_width=1, + sampler=Mock(), + waiting_queue=object(), + active_requests=[], + _mm_encoder_item_scheduling_enabled=False, + _fetch_new_requests=Mock(return_value=requests), + _configure_csa2_remote_tail=Mock(), + _validate_token_id_range=Mock(), + _validate_request_budget=Mock(), + _handle_errors=Mock(), + ) + executor._validate_request = types.MethodType(PyExecutor._validate_request, executor) + return executor + + +def _make_admission_request(request_id: int) -> types.SimpleNamespace: + return types.SimpleNamespace( + request_id=request_id, + py_disaggregated_params=None, + sampling_config=None, + kv_cache_retention_config=None, + ) + + +def test_model_validation_fails_only_the_request_it_rejects() -> None: + requests = [_make_admission_request(i) for i in range(3)] + + def validate_request(request) -> None: + if request.request_id == 1: + raise ValueError("frame too large") + + model = types.SimpleNamespace(validate_request=Mock(side_effect=validate_request)) + executor = _make_admission_executor(model, requests) + + admitted = PyExecutor._fetch_and_activate_new_requests(executor) + + assert admitted == [requests[0], requests[2]] + assert executor.active_requests == [requests[0], requests[2]] + executor._handle_errors.assert_called_once_with( + "frame too large", requests=[requests[1]], charge_budget=False + ) + assert [call.args[0] for call in model.validate_request.call_args_list] == requests + # The checks after the model's still run for the requests it admitted. + assert executor._validate_request_budget.call_count == 2 + + +def test_model_validation_runs_after_the_sampler_and_before_the_budget() -> None: + order = [] + model = types.SimpleNamespace(validate_request=lambda request: order.append("model")) + executor = _make_admission_executor(model, []) + executor.sampler.validate_request.side_effect = lambda request: order.append("sampler") + executor._validate_request_budget.side_effect = lambda request: order.append("budget") + + PyExecutor._validate_request(executor, _make_admission_request(0)) + + assert order == ["sampler", "model", "budget"] + + +@pytest.mark.parametrize( + "model", + [ + types.SimpleNamespace(), + Mock(spec=DecoderModelForCausalLM), + types.SimpleNamespace(validate_request=None), + None, + ], + ids=["no_attribute", "spec_mock", "none_attribute", "no_model"], +) +def test_models_without_the_validation_hook_are_unaffected(model) -> None: + requests = [_make_admission_request(i) for i in range(2)] + executor = _make_admission_executor(model, requests) + + admitted = PyExecutor._fetch_and_activate_new_requests(executor) + + assert admitted == requests + assert executor.active_requests == requests + executor._handle_errors.assert_not_called() + + @pytest.mark.parametrize("mode", ["standalone", "generation", "disabled"]) def test_full_model_requests_without_remote_handoff_are_unchanged(mode: str) -> None: executor, request = _make_remote_tail_executor_and_request(context_only=False) diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 3a7293bd7e9d..6e546e3a9a38 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -59,6 +59,7 @@ from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata from tensorrt_llm._torch.attention.backends.trtllm import \ TrtllmAttentionMetadata +from tensorrt_llm._torch.attention.backends.vanilla import VanillaAttention from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._torch.speculative.interface import \ INVALID_PROMPT_LOOKAHEAD_TOKEN @@ -67,7 +68,8 @@ from tensorrt_llm._torch.utils import set_per_request_prefill_cuda_graph_flag from tensorrt_llm.inputs.registry import (BaseMultimodalDummyInputsBuilder, BaseMultimodalInputProcessor) -from tensorrt_llm.llmapi import (CudaGraphConfig, SADecodingConfig, +from tensorrt_llm.llmapi import (CudaGraphConfig, MTPDecodingConfig, + NGramDecodingConfig, SADecodingConfig, SamplingParams) from tensorrt_llm.llmapi.llm_args import KvCacheConfig from tensorrt_llm.mapping import CpType, Mapping @@ -2898,11 +2900,15 @@ def test_prepare_inputs_with_helix_parallelism(self) -> None: actual_seq_lens = attn_metadata.seq_lens.cpu().tolist() self.assertEqual(actual_seq_lens, expected_seq_lens) - def _setup_mrope_engine(self, max_num_tokens: int = 32): + def _setup_mrope_engine(self, + max_num_tokens: int = 32, + spec_config: DecodingBaseConfig | None = None): """Build a DummyModelEngine that takes the MRoPE path, plus its KV cache manager and attention metadata.""" llm_args = TorchLlmArgs(model="dummy") - model_engine = DummyModelEngine(llm_args, dtype=torch.half) + model_engine = DummyModelEngine(llm_args, + dtype=torch.half, + spec_config=spec_config) model_engine.model.model_config.pretrained_config.rope_scaling = { "type": "mrope" } @@ -2931,6 +2937,393 @@ def _setup_mrope_engine(self, max_num_tokens: int = 32): _resize_input_buffers(model_engine._runner, max_num_tokens) return model_engine, kv_cache_manager, attn_metadata + @staticmethod + def _make_mrope_gen_request(num_tokens: int, req_id: int, seq_slot: int, + delta): + """Generation request carrying an MRoPE delta when `delta` is not None, + and no multimodal data at all otherwise (a text-only prompt).""" + request = _create_request(num_tokens, req_id) + request.py_prompt_len = num_tokens + request.py_batch_idx = None + request.py_seq_slot = seq_slot + request.sampling_config.beam_width = 1 + if delta is None: + request.py_multimodal_data = {} + else: + request.py_multimodal_data = { + "mrope_config": { + "mrope_position_deltas": + torch.tensor([[delta]], dtype=torch.int32) + }, + } + return request + + def _make_mrope_spec_batch(self, previous_batch: bool): + """Two multimodal verify groups around a text-only one, plus a CUDA graph + padding dummy, each carrying one draft token. + + Returns the scheduled batch and the multimodal requests.""" + # (num_tokens, req_id, seq_slot, delta); the middle request is text-only. + requests = [ + self._make_mrope_gen_request(4, 1, 0, 10), + self._make_mrope_gen_request(5, 2, 1, None), + self._make_mrope_gen_request(6, 3, 2, 20), + ] + for request in requests: + request.py_draft_tokens = [7] + if previous_batch: + # The overlap scheduler's previous step ran these requests at + # their current slots, which also seeded the delta cache. + request.py_batch_idx = request.py_seq_slot + if request.py_multimodal_data: + request.py_mrope_delta_cache_slot = request.py_seq_slot + # The padding dummy carries the zero delta that `add_dummy_requests` + # attaches on MRoPE models. + dummy = self._make_mrope_gen_request(2, 4, 3, 0) + dummy.is_cuda_graph_dummy = True + dummy.py_draft_tokens = [0] + + scheduled_requests = ScheduledRequests() + scheduled_requests.context_requests_last_chunk = [] + scheduled_requests.generation_requests = requests + [dummy] + return scheduled_requests, [requests[0], requests[2]] + + def test_prepare_inputs_spec_decode_adds_mrope_delta(self) -> None: + """One-model speculative verify groups get the request's MRoPE delta. + + Every token a generation step feeds is text after the prompt's vision + spans, so its three coordinates are `seq_index + delta`, as for an + ordinary generation request. Read slots stay dense, one per verify + group, and a generation-only request seeds the delta cache. + """ + model_engine, kv_cache_manager, attn_metadata = self._setup_mrope_engine( + spec_config=SADecodingConfig(max_draft_len=1)) + scheduled_requests, multimodal_requests = self._make_mrope_spec_batch( + previous_batch=False) + spec_metadata = Mock( + _force_non_greedy_for_capture=False, + context_prompt_lookahead_tokens=None, + ) + + result = model_engine._runner._inputs.prepare_inputs( + scheduled_requests=scheduled_requests, + kv_cache_manager=kv_cache_manager, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + enable_spec_decode=True, + runtime_draft_len=1, + is_dummy=False).kwargs + + # Each verify group is [last token, draft token] at past_seen + 0/1; + # the text-only request and the zero-delta dummy keep the scalar + # positions. + position_ids = result["position_ids"] + self.assertEqual(tuple(position_ids.shape), (3, 1, 8)) + expected = torch.tensor([[[13, 14, 4, 5, 25, 26, 1, 2]]] * 3, + dtype=torch.int32, + device='cuda') + torch.testing.assert_close(position_ids, expected, atol=0, rtol=0) + # Slot 32 is the reserved zero slot (max_num_tokens * pp_size). + self.assertEqual(result["mrope_delta_read_seq_slots"].cpu().tolist(), + [0, 32, 2, 32]) + self.assertEqual(result["mrope_delta_write_seq_slots"].cpu().tolist(), + [0, 2]) + self.assertEqual([ + request.py_mrope_delta_cache_slot for request in multimodal_requests + ], [0, 2]) + kv_cache_manager.shutdown() + + def test_prepare_inputs_spec_decode_overlap_keeps_mrope_delta(self) -> None: + """Under the overlap scheduler the accepted-token correction applied in + `preprocess_inputs` lands on top of the delta, on all three axes.""" + model_engine, kv_cache_manager, attn_metadata = self._setup_mrope_engine( + spec_config=SADecodingConfig(max_draft_len=1)) + scheduled_requests, _ = self._make_mrope_spec_batch(previous_batch=True) + spec_metadata = Mock( + _force_non_greedy_for_capture=False, + context_prompt_lookahead_tokens=None, + ) + # The previous step accepted 2, 1 and 2 tokens in slots 0, 1 and 2. + overlap_state = SampleStateTensorsSpec( + new_tokens=torch.zeros((2, 4, 1), dtype=torch.int32, device="cuda"), + new_tokens_lens=torch.tensor([2, 1, 2, 1], + dtype=torch.int32, + device="cuda"), + next_draft_tokens=torch.zeros((4, 1), + dtype=torch.int32, + device="cuda"), + ) + + inputs = model_engine._runner._inputs.prepare_inputs( + scheduled_requests=scheduled_requests, + kv_cache_manager=kv_cache_manager, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + new_tensors_device=overlap_state, + enable_spec_decode=True, + runtime_draft_len=1, + is_dummy=False).kwargs + + # The host packs positions from the stale base (past_seen + delta). + stale = torch.tensor([[[13, 14, 4, 5, 25, 26, 1, 2]]] * 3, + dtype=torch.int32, + device='cuda') + torch.testing.assert_close(inputs["position_ids"], + stale, + atol=0, + rtol=0) + self.assertEqual(inputs["mrope_delta_read_seq_slots"].cpu().tolist(), + [0, 32, 2, 32]) + # The delta cache already holds both multimodal requests' deltas. + self.assertNotIn("mrope_delta_write_seq_slots", inputs) + + inputs = model_engine._runner._inputs.preprocess_inputs( + inputs, enable_spec_decode=True, runtime_draft_len=1) + corrected = torch.tensor([[[15, 16, 5, 6, 27, 28, 1, 2]]] * 3, + dtype=torch.int32, + device='cuda') + torch.testing.assert_close(inputs["position_ids"], + corrected, + atol=0, + rtol=0) + kv_cache_manager.shutdown() + + @staticmethod + def _make_mrope_ctx_request(num_tokens: int, req_id: int, seq_slot: int, + delta: int): + """Single-chunk context request carrying MRoPE metadata. + + Its position on axis `a` of token `i` is `100 + 1000 * a + i`, so a + spliced span says which axis and token it came from.""" + request = _create_request(num_tokens, req_id) + request.py_prompt_len = num_tokens + request.py_batch_idx = None + request.py_seq_slot = seq_slot + request.context_current_position = 0 + request.context_chunk_size = num_tokens + request.sampling_config.beam_width = 1 + positions = (100 + 1000 * torch.arange(3).view(3, 1, 1) + + torch.arange(num_tokens).view(1, 1, -1)).to(torch.int32) + request.py_multimodal_data = { + "mrope_config": { + "mrope_position_ids": positions, + "mrope_position_deltas": torch.tensor([[delta]], + dtype=torch.int32), + }, + } + return request + + @staticmethod + def _prepare_mrope_spec_inputs(model_engine, kv_cache_manager, + attn_metadata, scheduled_requests): + spec_metadata = Mock( + _force_non_greedy_for_capture=False, + context_prompt_lookahead_tokens=None, + ) + result = model_engine._runner._inputs.prepare_inputs( + scheduled_requests=scheduled_requests, + kv_cache_manager=kv_cache_manager, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + enable_spec_decode=True, + runtime_draft_len=1, + is_dummy=False).kwargs + return result + + def _assert_mrope_positions(self, result, rows) -> None: + """`rows` is one list for all three axes, or one list per axis.""" + if not isinstance(rows[0], list): + rows = [rows] * 3 + position_ids = result["position_ids"] + self.assertEqual(tuple(position_ids.shape), (3, 1, len(rows[0]))) + expected = torch.tensor([rows[0:1], rows[1:2], rows[2:3]], + dtype=torch.int32, + device='cuda') + torch.testing.assert_close(position_ids, expected, atol=0, rtol=0) + + def test_prepare_inputs_spec_decode_mrope_plain_rows_follow_verify_groups( + self) -> None: + """Verify groups are packed ahead of plain generation requests, and the + MRoPE delta read slots follow that packing rather than scheduler order. + + The attention kernel indexes the deltas by generation batch index: row + i of the packed generation batch must read the delta of the request + packed i-th, so a verify group's slot comes before every plain + request's, whatever order the scheduler listed them in. + """ + model_engine, kv_cache_manager, attn_metadata = self._setup_mrope_engine( + spec_config=MTPDecodingConfig(max_draft_len=1)) + # `verify_*` carry a draft token, `plain_*` do not; `*_mm` have an + # MRoPE delta, `*_text` none. + plain_mm = self._make_mrope_gen_request(6, 1, 0, 20) + verify_mm = self._make_mrope_gen_request(4, 2, 1, 10) + plain_text = self._make_mrope_gen_request(5, 3, 2, None) + verify_text = self._make_mrope_gen_request(3, 4, 3, None) + verify_mm.py_draft_tokens = [7] + verify_text.py_draft_tokens = [7] + + scheduled_requests = ScheduledRequests() + scheduled_requests.context_requests_last_chunk = [] + scheduled_requests.generation_requests = [ + plain_mm, verify_mm, plain_text, verify_text + ] + result = self._prepare_mrope_spec_inputs(model_engine, kv_cache_manager, + attn_metadata, + scheduled_requests) + + # Packed order: verify_mm, verify_text, plain_mm, plain_text. The + # verify groups hold two rows each, the plain requests one. + self._assert_mrope_positions(result, [13, 14, 2, 3, 25, 4]) + # Slot 32 is the reserved zero slot (max_num_tokens * pp_size) read by + # the text-only requests. + self.assertEqual(result["mrope_delta_read_seq_slots"].cpu().tolist(), + [1, 32, 0, 32]) + self.assertEqual(result["mrope_delta_write_seq_slots"].cpu().tolist(), + [1, 0]) + kv_cache_manager.shutdown() + + def test_prepare_inputs_spec_decode_mrope_context_takes_no_read_slot( + self) -> None: + """A context request packs ahead of the verify groups with its own + MRoPE positions and takes no generation read slot, so the verify + groups' slots start the list.""" + model_engine, kv_cache_manager, attn_metadata = self._setup_mrope_engine( + spec_config=MTPDecodingConfig(max_draft_len=1)) + # A context request with multimodal content takes the placeholder-token + # split, which reads the vocabulary size. + model_engine.model.config.vocab_size = 1000 + context = self._make_mrope_ctx_request(5, 1, 0, 30) + verify_mm = self._make_mrope_gen_request(4, 2, 1, 10) + verify_mm.py_draft_tokens = [7] + plain_mm = self._make_mrope_gen_request(6, 3, 2, 20) + + scheduled_requests = ScheduledRequests() + scheduled_requests.context_requests_last_chunk = [context] + scheduled_requests.generation_requests = [plain_mm, verify_mm] + result = self._prepare_mrope_spec_inputs(model_engine, kv_cache_manager, + attn_metadata, + scheduled_requests) + + # The context span keeps the processor's per-axis positions; the + # generation rows after it are text at position + delta on all axes. + self._assert_mrope_positions(result, [ + [100, 101, 102, 103, 104, 13, 14, 25], + [1100, 1101, 1102, 1103, 1104, 13, 14, 25], + [2100, 2101, 2102, 2103, 2104, 13, 14, 25], + ]) + self.assertEqual(result["mrope_delta_read_seq_slots"].cpu().tolist(), + [1, 2]) + # The context request seeds the delta cache first, then the verify + # group and the plain request seed theirs. + self.assertEqual(result["mrope_delta_write_seq_slots"].cpu().tolist(), + [0, 1, 2]) + self.assertEqual(context.py_mrope_delta_cache_slot, 0) + kv_cache_manager.shutdown() + + def _setup_extend_as_context_engine(self): + """A two-model draft mode on an attention backend that runs verify groups + as chunked context rows. + + `extend_ctx` is False for every one-model mode and for the TRTLLM + backend below sm_120, so this takes a drafter-based mode and swaps in + the vanilla backend. + """ + model_engine, kv_cache_manager, attn_metadata = self._setup_mrope_engine( + spec_config=NGramDecodingConfig(max_draft_len=1)) + _set_runner_config(model_engine._runner, + attention_backend=VanillaAttention, + disable_overlap_scheduler=True) + spec_config = model_engine._runner._config.spec_config + self.assertTrue(spec_config.spec_dec_mode.extend_ctx(VanillaAttention)) + self.assertTrue(spec_config.is_linear_tree) + return model_engine, kv_cache_manager, attn_metadata + + def test_prepare_inputs_mrope_verify_as_context_takes_no_read_slot( + self) -> None: + """When the kernel runs verify groups as chunked context rows, their + rotary comes from `position_ids` alone and only the plain generation + requests take read slots.""" + model_engine, kv_cache_manager, attn_metadata = ( + self._setup_extend_as_context_engine()) + verify_mm = self._make_mrope_gen_request(4, 1, 0, 10) + verify_text = self._make_mrope_gen_request(3, 2, 1, None) + plain_mm = self._make_mrope_gen_request(6, 3, 2, 20) + plain_text = self._make_mrope_gen_request(5, 4, 3, None) + verify_mm.py_draft_tokens = [7] + verify_text.py_draft_tokens = [7] + + scheduled_requests = ScheduledRequests() + scheduled_requests.context_requests_last_chunk = [] + scheduled_requests.generation_requests = [ + verify_mm, verify_text, plain_mm, plain_text + ] + result = self._prepare_mrope_spec_inputs(model_engine, kv_cache_manager, + attn_metadata, + scheduled_requests) + + self._assert_mrope_positions(result, [13, 14, 2, 3, 25, 4]) + self.assertEqual(result["mrope_delta_read_seq_slots"].cpu().tolist(), + [2, 32]) + # The verify group still seeds the delta cache for the plain steps + # that follow it. + self.assertEqual(result["mrope_delta_write_seq_slots"].cpu().tolist(), + [0, 2]) + kv_cache_manager.shutdown() + + def test_prepare_inputs_mrope_verify_as_context_alone_emits_no_read_slots( + self) -> None: + """A batch whose only delta carriers are verify groups run as context + rows has nothing for the generation-indexed read.""" + model_engine, kv_cache_manager, attn_metadata = ( + self._setup_extend_as_context_engine()) + verify_mm = self._make_mrope_gen_request(4, 1, 0, 10) + verify_mm.py_draft_tokens = [7] + plain_text = self._make_mrope_gen_request(5, 2, 1, None) + + scheduled_requests = ScheduledRequests() + scheduled_requests.context_requests_last_chunk = [] + scheduled_requests.generation_requests = [verify_mm, plain_text] + result = self._prepare_mrope_spec_inputs(model_engine, kv_cache_manager, + attn_metadata, + scheduled_requests) + + self._assert_mrope_positions(result, [13, 14, 4]) + self.assertNotIn("mrope_delta_read_seq_slots", result) + self.assertEqual(result["mrope_delta_write_seq_slots"].cpu().tolist(), + [0]) + kv_cache_manager.shutdown() + + def _make_first_draft_batch(self): + """A first-draft request ahead of two plain generation requests that + carry MRoPE deltas.""" + first_draft = self._make_mrope_gen_request(4, 1, 0, 10) + first_draft.py_is_first_draft = True + first_draft.py_num_accepted_draft_tokens = 0 + scheduled_requests = ScheduledRequests() + scheduled_requests.context_requests_last_chunk = [] + scheduled_requests.generation_requests = [ + first_draft, + self._make_mrope_gen_request(6, 2, 1, 20), + self._make_mrope_gen_request(5, 3, 2, 30), + ] + return scheduled_requests + + def test_prepare_inputs_first_draft_requests_without_mrope_are_served( + self) -> None: + """Without MRoPE a first-draft request packs as a plain generation row.""" + model_engine, kv_cache_manager, attn_metadata = self._setup_mrope_engine( + spec_config=MTPDecodingConfig(max_draft_len=1)) + _set_runner_config(model_engine._runner, use_mrope=False) + result = self._prepare_mrope_spec_inputs(model_engine, kv_cache_manager, + attn_metadata, + self._make_first_draft_batch()) + + # The first-draft row is [last two prompt tokens + draft slot] wide. + self.assertEqual(result["position_ids"].view(-1).cpu().tolist(), + [2, 3, 5, 4]) + self.assertNotIn("mrope_delta_read_seq_slots", result) + kv_cache_manager.shutdown() + @skip_pre_blackwell_unittest def test_ced_warmup_passes_dummy_context_requests(self) -> None: for warmup in (False, True): diff --git a/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py b/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py index ee5b97703635..6d8d1b727a55 100644 --- a/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py +++ b/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py @@ -365,7 +365,7 @@ def test_qsa_cache_manager_requires_sparse_config() -> None: QSAMambaHybridCacheManagerV2(layer_mask=[True]) -@pytest.mark.parametrize("dtype", (DataType.NVFP4, DataType.FLOAT)) +@pytest.mark.parametrize("dtype", (DataType.FLOAT,)) def test_qsa_cache_manager_delegates_regular_kv_layout( monkeypatch: pytest.MonkeyPatch, dtype: DataType ) -> None: diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py index b3e446dd7b23..5c196b997868 100644 --- a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py +++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py @@ -175,7 +175,8 @@ within-block mask read out one-hot at P = 2, 3, 4 over a bf16 pool at 24/2/256 and 24/4/256 and an fp8 one at 24/4/256 (bottom-right causal, later siblings bitwise zero, mask_type inert, a page-crossing block, - consecutive steps), with the full-block model as the measured rival. + consecutive steps), with the full-block model as the measured rival; the + entry's cells gate P up to 8 (the target's max_draft_len 7) against `reference`. 12. qwen4_exp's full-attention cell, GQA 24/2/256 over a bf16 pool on the paged-context path: the cached-prefix rival, the sparse / skip-softmax group measured bitwise inert, and local_layer_idx over its 48-row mapping @@ -183,11 +184,16 @@ real 48-layer manager's identity rows, and a bitwise A/B of a single-layer pool against mapping row 47. 13. The no-KV-cache (encoder-only) mode qwen4_exp's vision tower runs: all - three cache-addressing tensors None, one context sequence per frame, MHA - 16/16 at head_size 72 and 128, the merged-segment rival, the cache + three cache-addressing tensors None, one context sequence per frame (an + image, or one frame group of a video), MHA 16/16 at head_size 72 and 128, + the merged-segment rival, the cache arguments that are inert there, the workspace high-water mark over the mmmu serving shapes, and the two boundaries (update_kv_cache=False is rejected; head_size 72 has no generation phase). +14. The same full-attention cell over an NVFP4 pool (e2m1 data and e4m3 block + scales in a second pool, unit fp32 [3] second-level scales, bf16 output): + the pool and its block scales bitwise as the append writes them, and the + output against `reference` over the NVFP4 values (the cells). The cells at the end drive the entry's own `reference` and `_compare` over the configurations the targets run. """ @@ -241,6 +247,7 @@ MASK_PADDING = 0 # AttentionMaskType.padding: no mask (bidirectional) QUANT_MODE_FP8_KV_CACHE = 128 # QuantMode.FP8_KV_CACHE (sole bit set) +QUANT_MODE_NVFP4_KV_CACHE = 8192 # QuantMode.NVFP4_KV_CACHE (sole bit set) # fp8-e4m3 KV-pool decode tolerance. The reference dequantizes the mirrored # e4m3 pool content and puts q through the same e4m3 round-trip the kernel @@ -291,6 +298,88 @@ FP8_MLA_RTOL = 2**-4 +# NVFP4: one e4m3 block scale per this many values; e2m1 magnitudes by code +# and the round-to-nearest-even steps between them (inclusive where the upper +# neighbour is the even code). +NVFP4_BLOCK = 16 +E2M1_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) +E2M1_STEPS = ((0.25, False), (0.75, True), (1.25, False), (1.75, True), (2.5, False)) +E2M1_STEPS += ((3.5, True), (5.0, False)) + + +def _nvfp4_quantize(rows: torch.Tensor, global_scale: torch.Tensor): + """[n, D] rows -> (packed e2m1 bytes [n, D/2], e4m3 scale bytes [n, D/16]): + the attention op's append, in torch (`cvt_warp_fp16_to_fp4`).""" + n, d = rows.shape + x = rows.float().reshape(n, d // NVFP4_BLOCK, NVFP4_BLOCK) + block_max = x.abs().amax(-1) + sixth = torch.tensor(6.0, dtype=torch.float32, device=rows.device).reciprocal() + code = (global_scale * (block_max * sixth)).to(torch.float8_e4m3fn) + out_scale = torch.where( + block_max != 0, + (code.float() * global_scale.reciprocal()).reciprocal(), + torch.zeros_like(block_max), + ) + scaled = x * out_scale[..., None] + magnitude = scaled.abs() + codes = torch.zeros(scaled.shape, dtype=torch.int32, device=rows.device) + for threshold, inclusive in E2M1_STEPS: + codes += (magnitude >= threshold if inclusive else magnitude > threshold).to(torch.int32) + codes |= torch.signbit(scaled).to(torch.int32) << 3 + # 0 x inf, in a block whose scale underflows to zero: the conversion makes +6. + codes = torch.where(torch.isnan(scaled), torch.full_like(codes, 0x7), codes) + codes = codes.reshape(n, d // 2, 2) + return (codes[..., 0] | (codes[..., 1] << 4)).to(torch.uint8), code.view(torch.uint8) + + +def _nvfp4_dequantize(packed: torch.Tensor, scales: torch.Tensor, global_scale) -> torch.Tensor: + """Inverse of `_nvfp4_quantize`: fp32 [n, D].""" + lut = torch.tensor(E2M1_VALUES, dtype=torch.float32, device=packed.device) + lut = torch.cat([lut, -lut]) + raw = packed.long() + values = torch.stack([lut[raw & 0xF], lut[raw >> 4]], -1).flatten(-2) + scale = scales.view(torch.float8_e4m3fn).float().repeat_interleave(NVFP4_BLOCK, -1) + return values * scale / global_scale + + +def _nvfp4_roundtrip(kv: torch.Tensor, global_scale) -> torch.Tensor: + """[T, Hkv, D] bf16 -> the bf16 values an NVFP4 pool returns for them.""" + t, h, d = kv.shape + packed, scales = _nvfp4_quantize(kv.reshape(t * h, d), global_scale) + return _nvfp4_dequantize(packed, scales, global_scale).reshape(t, h, d).to(kv.dtype) + + +def _check_nvfp4_cache(env, request_id: int) -> None: + """An NVFP4 pool holds each token's e2m1 bytes and its block scales -- K's + token-major, V's in four-token groups -- bitwise as the append writes them.""" + tpb, hkv, d = env.tokens_per_block, env.num_kv_heads, env.head_dim + groups = d // NVFP4_BLOCK + for role, history in ((0, env.k_history[request_id]), (1, env.v_history[request_id])): + seq = torch.cat(history) + total = seq.shape[0] + packed, scales = _nvfp4_quantize( + seq.reshape(total * hkv, d), env.kv_scale_orig_quant[1 + role] + ) + packed = packed.reshape(total, hkv, d // 2) + scales = scales.reshape(total, hkv, groups) + for j, page in enumerate(env.pages[request_id]): + n = min(tpb, total - j * tpb) + if n <= 0: + break + rows = slice(j * tpb, j * tpb + n) + got = env.kv_cache[page, role, :, :n].permute(1, 0, 2) + assert torch.equal(got, packed[rows]), f"request {request_id} role {role} page {j}" + sf_page = env.kv_block_scales[page, role] + if role == 0: + got_sf = sf_page[:, :n].permute(1, 0, 2) + else: + grouped = sf_page.reshape(hkv, tpb // 4, groups, 4).permute(0, 1, 3, 2) + got_sf = grouped.reshape(hkv, tpb, groups)[:, :n].permute(1, 0, 2) + assert torch.equal(got_sf, scales[rows]), ( + f"request {request_id} role {role} page {j}: block scales" + ) + + def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: """A softmax chain, so the band is the longest in this catalog. @@ -378,26 +467,46 @@ def __init__( self.page_ring = page_ring self.rings: dict[int, List[int]] = {} - # Single-layer, single-pool paged KV cache (HND layout). + # Single-layer, single-pool paged KV cache (HND layout). An NVFP4 pool + # holds e2m1 pairs, and its e4m3 block scales live in a second pool + # whose pointer is the third axis of the pool pointers. + self.nvfp4 = bool(quant_mode & QUANT_MODE_NVFP4_KV_CACHE) self.kv_cache = torch.zeros( num_blocks, 2, num_kv_heads, tokens_per_block, - head_dim, - dtype=pool_dtype, + head_dim // 2 if self.nvfp4 else head_dim, + dtype=torch.uint8 if self.nvfp4 else pool_dtype, device="cuda", ) - self.pool_pointers = torch.zeros(1, 2, dtype=torch.int64, device="cpu") - self.pool_pointers[0, 0] = self.kv_cache.data_ptr() + if self.nvfp4: + self.kv_block_scales = torch.zeros( + num_blocks, + 2, + num_kv_heads, + tokens_per_block, + head_dim // NVFP4_BLOCK, + dtype=torch.uint8, + device="cuda", + ) + self.pool_pointers = torch.zeros(1, 2, 2, dtype=torch.int64, device="cpu") + self.pool_pointers[0, 0, 0] = self.kv_cache.data_ptr() + self.pool_pointers[0, 0, 1] = self.kv_block_scales.data_ptr() + else: + self.pool_pointers = torch.zeros(1, 2, dtype=torch.int64, device="cpu") + self.pool_pointers[0, 0] = self.kv_cache.data_ptr() # KV-cache scaling factor s: dequant = quantized * s. Mirrors the # production construction (quant_orig = s, orig_quant = 1/s, both - # fp32 [1] CUDA); None = unquantized pool, scale args passed as None. + # fp32 [1] CUDA, [3] -- q, K, V -- over an NVFP4 pool); None = + # unquantized pool, scale args passed as None. if kv_scaling_factor is None: self.kv_scale_quant_orig: Optional[torch.Tensor] = None self.kv_scale_orig_quant: Optional[torch.Tensor] = None else: - factor = torch.full((1,), kv_scaling_factor, dtype=torch.float32, device="cuda") + factor = torch.full( + (3 if self.nvfp4 else 1,), kv_scaling_factor, dtype=torch.float32, device="cuda" + ) self.kv_scale_quant_orig = factor self.kv_scale_orig_quant = 1.0 / factor self.pool_mapping = torch.zeros(1, 2, dtype=torch.int32, device="cpu") @@ -702,7 +811,11 @@ def cache_pages_content(self, request_id: int) -> Tuple[torch.Tensor, ...]: def check_cache(self, request_id: int) -> None: """The paged cache must hold exactly the expected rows, bitwise (the - bf16 K/V fed so far, or their scaled e4m3 casts for an fp8 pool).""" + bf16 K/V fed so far, or their scaled e4m3 casts for an fp8 pool, or + their NVFP4 bytes and block scales for an NVFP4 pool).""" + if self.nvfp4: + _check_nvfp4_cache(self, request_id) + return got_k, got_v, exp_k, exp_v = self.cache_pages_content(request_id) assert torch.equal(got_k.view(torch.uint8), exp_k.view(torch.uint8)), ( f"K cache mismatch: request {request_id}" @@ -5426,9 +5539,13 @@ def _op_reference( causal: bool = True, e4m3: bool = False, layer: Optional[int] = None, + nvfp4: bool = False, ) -> torch.Tensor: """`op.reference` over the history the env mirrored for these requests. + nvfp4=True reads K and V as an NVFP4 pool returns them (the history + through `_nvfp4_roundtrip`) before the e4m3 arithmetic `e4m3` selects. + Call it after the op call, so each history ends with this call's K/V: `reference`'s bottom-right alignment then puts this call's q rows at the end of it. causal=False on a generation call is the full-block rival -- @@ -5438,6 +5555,9 @@ def _op_reference( """ hq, hkv, d = env.num_heads, env.num_kv_heads, env.head_dim ks, vs = zip(*(_history(env, rid, layer) for rid in request_ids)) + if nvfp4: + ks = tuple(_nvfp4_roundtrip(k, env.kv_scale_orig_quant[1]) for k in ks) + vs = tuple(_nvfp4_roundtrip(v, env.kv_scale_orig_quant[2]) for v in vs) return op.reference( qkv[:, : hq * d], torch.cat(ks).reshape(-1, hkv * d), @@ -5930,8 +6050,8 @@ def repeat_pools(alias: bool) -> List[torch.Tensor]: # (its draft chain), token-major. On sm_103 a linear draft leaves every # spec_decoding_* tensor None and is_spec_decoding_enabled False, so the # within-block mask has to come from predicted_tokens_per_seq alone. The cells -# gate P = 1..4 against `op.reference`; these read the mask out directly and -# measure the rival. +# gate P = 1..8 against `op.reference`; these read the mask out directly and +# measure the rival at P = 2..4. # One configuration per pool the cells certify P over: qwen4_exp's 24/2/256 # over bf16, and 24/4/256 over bf16 and over fp8 at s = 1.0. The weight @@ -5942,6 +6062,7 @@ def repeat_pools(alias: bool) -> List[torch.Tensor]: ("bf16 24/2/256", 2, torch.bfloat16, 0, None, SWA_WEIGHT_RTOL), ("bf16 24/4/256", 4, torch.bfloat16, 0, None, SWA_WEIGHT_RTOL), ("fp8 24/4/256", 4, torch.float8_e4m3fn, QUANT_MODE_FP8_KV_CACHE, 1.0, 2**-4), + ("fp8 24/2/256", 2, torch.float8_e4m3fn, QUANT_MODE_FP8_KV_CACHE, 1.0, 2**-4), ] @@ -6261,7 +6382,11 @@ def again(mask_type: int = MASK_CAUSAL, **overrides) -> torch.Tensor: def _qwen_layer_pools( - attn_blocks: int = 48, other_blocks: int = 32, max_seq_len: int = 256, max_batch: int = 4 + attn_blocks: int = 48, + other_blocks: int = 32, + max_seq_len: int = 256, + max_batch: int = 4, + **pool_kw, ) -> Tuple[_MultiLayerPagedAttnEnv, _MultiLayerPagedAttnEnv]: """The two pools of the hybrid mapping, at the full-attention cell. @@ -6269,7 +6394,8 @@ def _qwen_layer_pools( rows name a real, distinct pool base and the untouched assertion has somewhere to look. Both managers take the same max_seq_len and max_batch, which makes their block-offset tables concatenable into the per-pool table - the op indexes by the mapping's pool column. + the op indexes by the mapping's pool column. `pool_kw` (dtype, quant_mode, + kv_scaling_factor) goes to both. """ kw = dict( num_heads=QWEN_HQ, @@ -6278,6 +6404,7 @@ def _qwen_layer_pools( tokens_per_block=QWEN_TPB, max_seq_len=max_seq_len, max_batch=max_batch, + **pool_kw, ) attn = _MultiLayerPagedAttnEnv(num_layers=len(QWEN_ATTN_LAYERS), num_blocks=attn_blocks, **kw) other = _MultiLayerPagedAttnEnv( @@ -6971,13 +7098,14 @@ def _cell_paged(cell, arm) -> None: A first, uncertified context call fills each sequence's `cached` tokens; the certified call then appends `new` rows per sequence (P for a generation row) and is gated against `op.reference` over the history it should have - attended. Over an fp8 pool every operand that call reads went through + attended. Over an fp8 or NVFP4 pool every operand that call reads went through e4m3, except on the packed context path, which never reads the pool. """ spec = cell.spec arm(op) _seed(cell) fp8 = bool(spec["quant_mode"] & QUANT_MODE_FP8_KV_CACHE) + nvfp4 = bool(spec["quant_mode"] & QUANT_MODE_NVFP4_KV_CACHE) env = _PagedAttnEnv( num_heads=spec["num_heads"], num_kv_heads=spec["num_kv_heads"], @@ -6988,7 +7116,7 @@ def _cell_paged(cell, arm) -> None: max_blocks_per_seq=16, pool_dtype=torch.float8_e4m3fn if fp8 else torch.bfloat16, quant_mode=spec["quant_mode"], - kv_scaling_factor=spec["kv_scale"] if fp8 else None, + kv_scaling_factor=spec["kv_scale"] if fp8 or nvfp4 else None, ) cached, new, nc = spec["cached"], spec["new"], spec["num_contexts"] p, paged = spec["predicted_tokens_per_seq"], spec["use_paged_context_fmha"] @@ -7023,9 +7151,10 @@ def _cell_paged(cell, arm) -> None: # One e4m3 flag covers the whole call: the packed path's context rows read # bf16, its generation rows e4m3, so a cell never mixes the two there. assert not (fp8 and not paged and 0 < nc < len(rids)) - e4m3 = fp8 and (paged or nc == 0) + assert not (nvfp4 and not paged), "an NVFP4 cell runs the paged-context path" + e4m3 = (fp8 or nvfp4) and (paged or nc == 0) causal = spec["mask_type"] == MASK_CAUSAL - _compare(out, _op_reference(env, qkv, list(new), rids, causal=causal, e4m3=e4m3)) + _compare(out, _op_reference(env, qkv, list(new), rids, causal=causal, e4m3=e4m3, nvfp4=nvfp4)) for rid in rids: env.check_cache(rid) env.check_unwritten_pool_zero() @@ -7046,8 +7175,14 @@ def _cell_layered(cell, arm) -> None: row, p = spec["local_layer_idx"], spec["predicted_tokens_per_seq"] cached, new, nc = spec["cached"], spec["new"], spec["num_contexts"] assert all(cached), "the history call drives every sequence of the cell" + fp8 = bool(spec["quant_mode"] & QUANT_MODE_FP8_KV_CACHE) + pool_kw = ( + dict(dtype=DataType.FP8, quant_mode=spec["quant_mode"], kv_scaling_factor=spec["kv_scale"]) + if fp8 + else {} + ) if spec["pool_layout"] == "hybrid48": - env, other = _qwen_layer_pools(max_batch=spec["max_num_requests"]) + env, other = _qwen_layer_pools(max_batch=spec["max_num_requests"], **pool_kw) layer = QWEN_ATTN_LAYERS.index(row) else: env = _MultiLayerPagedAttnEnv( @@ -7059,6 +7194,7 @@ def _cell_layered(cell, arm) -> None: num_blocks=24, max_batch=spec["max_num_requests"], max_seq_len=128, + **pool_kw, ) other, layer = None, row @@ -7104,7 +7240,7 @@ def tables() -> Dict: predicted_tokens_per_seq=p, **tables(), ) - _compare(out, _op_reference(env, qkv, list(new), rids, layer=layer)) + _compare(out, _op_reference(env, qkv, list(new), rids, layer=layer, e4m3=fp8)) env.check_caches(rids, layers=[layer]) _assert_only_this_layer_moved(env, snapshot, layer, cell.why) for before, after in zip(other_before, [] if other is None else other.layer_views): diff --git a/tests/unittest/_torch/modeling_v2/graph/test_modeling_v2_eager_on_graph.py b/tests/unittest/_torch/modeling_v2/graph/test_modeling_v2_eager_on_graph.py new file mode 100644 index 000000000000..33e34b6db83b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/graph/test_modeling_v2_eager_on_graph.py @@ -0,0 +1,285 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the eager_on_graph catalog entry. + +Each cell captures (or, for a passthrough cell, eagerly runs) a body of `breaks` breaks with +row-wise ops between them, the shape of a breakable prefill body, and replays it at the cell's +live rows. Every break's `inner` reads the step's rows and scale from live state, mixes rows +`[:n]` of its operand across the batch, writes a fixed output allocated inside the body and +zeroes its rows `[n:]`. The replayed body must equal the same body run eagerly at the same +live state, bit for bit; the hazards the entry's note states are each shown by a test below. +""" + +import ctypes +import logging +import zlib +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.graph.eager_on_graph import ( + eager_on_graph as op, +) +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import ( + BreakableCUDAGraph, + BreakableCUDAGraphCapture, +) + +__extra_import_path__ = [".."] +from _validating import certifying # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "eager_on_graph requires a CUDA device" + +_BCG_LOGGER = "tensorrt_llm._torch.pyexecutor.breakable_cuda_graph.breakable_cuda_graph" + +#: `cudaError_t` values the pending-error checks compare against. +_CUDA_SUCCESS = 0 +_CUDA_ERROR_ILLEGAL_STATE = 401 + + +def _loaded_cudart() -> ctypes.CDLL: + """The CUDA runtime library torch already loaded into this process. Its + per-thread last error is the one torch's launch checks read, which + `cuda.bindings.runtime` (a separate runtime instance) does not see.""" + soname = f"/libcudart.so.{torch.version.cuda.split('.')[0]}" + with open("/proc/self/maps") as maps: + paths = sorted({line.split()[-1] for line in maps if soname in line}) + assert len(paths) == 1, f"expected exactly one loaded {soname[1:]}, found {paths}" + return ctypes.CDLL(paths[0]) + + +_CUDART = _loaded_cudart() + + +def _take_pending_cuda_error() -> int: + """Read and clear the CUDA runtime's pending last error on this thread.""" + return int(_CUDART.cudaGetLastError()) + + +@pytest.fixture(autouse=True) +def _leaves_cuda_clean(): + """Every test here hands the process on as it found it: no stream + capturing and no CUDA runtime error pending. A pending error is reported + by whichever CUDA launch comes next -- in a shared pytest process, the + first call of the next test file -- so it is taken here, after every test + and whether or not the test passed, and a test that left one fails here + instead of there.""" + yield + pending = _take_pending_cuda_error() + torch.cuda.synchronize() + assert not torch.cuda.is_current_stream_capturing(), "a test left the stream capturing" + assert pending == _CUDA_SUCCESS, f"the test left CUDA runtime error {pending} pending" + + +class _Live: + """The per-step state a break reads, refreshed before every step.""" + + def __init__(self, rows: int, scale: float) -> None: + self.rows = rows + self.scale = scale + + +def _inner(live: _Live): + """A break's eager work: rows `[:n]` centred over the live batch, scaled by + the live scale and offset by the static layer index, into `out` in place; + rows `[n:]` zeroed. Returns None, as an in-place break does.""" + + def inner(src: torch.Tensor, out: torch.Tensor, layer: int) -> None: + n = live.rows + rows = src[:n].float() + mixed = (rows - rows.mean(0, keepdim=True)) * live.scale + layer + out[:n].copy_(mixed.to(out.dtype)) + out[n:].zero_() + + return inner + + +def _body(x: torch.Tensor, breaks: int, call) -> tuple[torch.Tensor, list[torch.Tensor]]: + """Row-wise ops around `breaks` breaks; each break's output is allocated in + the body, as a captured body allocates it from the graph pool.""" + h = torch.sin(x) * 2.0 + outs = [] + for layer in range(breaks): + out = torch.empty_like(h) + call(h, out, layer) + outs.append(out) + h = torch.tanh(out) + h * 0.5 + return h, outs + + +def _capture(x_static: torch.Tensor, body) -> tuple: + """Two eager warmups, then the capture, on a side stream as the engine runs it.""" + graph = BreakableCUDAGraph() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(2): + body(x_static) + with BreakableCUDAGraphCapture(graph, pool=torch.cuda.graph_pool_handle(), stream=stream): + h, outs = body(x_static) + torch.cuda.current_stream().wait_stream(stream) + return graph, h, outs + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + spec = cell.spec + tokens, width, breaks = spec["tokens"], spec["width"], spec["breaks"] + gen = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + live = _Live(rows=tokens, scale=1.0) + inner = _inner(live) + + def broken(*operands): + return op(*operands, inner=inner) + + def eager(*operands): + return op.reference(*operands, inner=inner) + + x_static = torch.randn([tokens, width], generator=gen, device="cuda").to(torch.bfloat16) + graph = None + if spec["capture"]: + graph, h_static, outs_static = _capture(x_static, lambda x: _body(x, breaks, broken)) + assert graph.num_breaks == breaks and graph.num_segments == breaks + 1 + assert graph.num_segments == breaks + 1 + + for step, rows in enumerate(spec["live_rows"]): + assert 0 < rows <= tokens + live.rows, live.scale = rows, 0.25 + 0.5 * step + x_new = torch.randn([tokens, width], generator=gen, device="cuda").to(torch.bfloat16) + if graph is not None: + x_static.copy_(x_new) + graph.replay() + got, got_outs = h_static.clone(), [o.clone() for o in outs_static] + else: + got, got_outs = _body(x_new, breaks, broken) + want, want_outs = _body(x_new, breaks, eager) + assert len(got_outs) == len(want_outs) == breaks + op.compare(got, want) + for got_out, want_out in zip(got_outs, want_outs): + op.compare(got_out, want_out) + assert torch.count_nonzero(got_out[rows:]) == 0, "padding rows left unwritten" + + +def _single_break_graph(inner, x: torch.Tensor, *extra): + """Capture `inner(h, out, *extra)` as one break between two segments that + read `x` and copy `out` into the returned result. Every tensor the graph + touches is allocated inside the capture and kept by the caller.""" + graph = BreakableCUDAGraph() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + with BreakableCUDAGraphCapture(graph, pool=torch.cuda.graph_pool_handle(), stream=stream): + h = x * 1.0 + out = torch.empty_like(h) + op(h, out, *extra, inner=inner) + result = out + 0.0 + torch.cuda.current_stream().wait_stream(stream) + return graph, result, (h, out) + + +def test_a_scalar_argument_is_frozen_at_capture() -> None: + """A per-step count handed in as an argument replays at its capture-time value.""" + + def inner(src, out, n): + out[:n].copy_(src[:n] * 2) + out[n:].zero_() + + x = torch.randn(64, 8, device="cuda") + graph, result, _keep = _single_break_graph(inner, x, 64) + # The step now carries 17 rows, but nothing can hand the replay that number. + graph.replay() + torch.cuda.synchronize() + assert torch.equal(result, x * 2), "the replay must run at the capture-time n" + stale = torch.zeros_like(x) + stale[:17] = x[:17] * 2 + assert not torch.equal(result, stale) + + +def test_a_tensor_argument_is_read_at_its_capture_address() -> None: + """A tensor argument is re-read at its capture-time address: a new tensor the + step computed is never seen, and the captured one is read as it now is.""" + + def inner(src, out, scale): + torch.mul(src, scale, out=out) + + x = torch.randn(32, 8, device="cuda") + captured = torch.full((1,), 1.0, device="cuda") + graph, result, _keep = _single_break_graph(inner, x, captured) + fresh = torch.full((1,), 2.0, device="cuda") # what a per-step computation would make + assert fresh.data_ptr() != captured.data_ptr() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(result, x * 1.0), "replay must read the capture-time tensor" + captured.fill_(3.0) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(result, x * 3.0), "replay must read the capture-time address as it is now" + + +def test_a_wait_on_a_non_capturing_stream_is_dropped(caplog) -> None: + """Inside the inner's own capture-time run, a wait on a stream that is not + capturing is dropped, with a warning naming the call site.""" + side = torch.cuda.Stream() + + def inner(src, out, layer): + torch.cuda.current_stream().wait_stream(side) + out.copy_(src) + + with caplog.at_level(logging.WARNING, logger=_BCG_LOGGER): + _single_break_graph(inner, torch.randn(8, 8, device="cuda"), 0) + assert any("Dropping a wait" in record.message for record in caplog.records) + + +def test_an_object_argument_raises_under_capture() -> None: + """Any argument the replay cannot hold weakly is refused at capture, after + the inner already ran once. Outside capture the same call is a passthrough. + The refusal is a break raising under capture, so it also leaves the error + `test_a_raising_break_leaves_an_error_pending` shows, taken here.""" + ran = [] + + def inner(src, out, config): + ran.append(config.tag) + out.copy_(src) + + config = SimpleNamespace(tag="step") + x = torch.randn(8, 8, device="cuda") + out = torch.empty_like(x) + op(x, out, config, inner=inner) + assert torch.equal(out, x) and ran == ["step"] + + with pytest.raises((TypeError, RuntimeError)) as info: + _single_break_graph(inner, x, config) + chain = [] + error = info.value + while error is not None: + chain.append(error) + error = error.__context__ + assert any(isinstance(e, TypeError) and "Invalid type" in str(e) for e in chain) + assert ran == ["step", "step"] + assert _take_pending_cuda_error() == _CUDA_ERROR_ILLEGAL_STATE + torch.cuda.synchronize() + + +def test_a_raising_break_leaves_an_error_pending() -> None: + """A break that raises under capture has already ended the captured + segment, and the capture's exit ends it again: the exit raises + RuntimeError over the break's own error and leaves cudaErrorIllegalState + pending, which the process's next CUDA launch reports as its own failure + -- once, since that launch's check clears it.""" + + def inner(src, out, layer): + out.copy_(src) + raise ValueError("the break's own failure") + + with pytest.raises(RuntimeError) as info: + _single_break_graph(inner, torch.randn(8, 8, device="cuda"), 0) + assert isinstance(info.value.__context__, ValueError), info.value.__context__ + torch.cuda.synchronize() + + assert _CUDART.cudaPeekAtLastError() == _CUDA_ERROR_ILLEGAL_STATE + with pytest.raises(torch.AcceleratorError, match="cannot be performed in the present state"): + torch.ones(1, device="cuda").add_(1) + assert _take_pending_cuda_error() == _CUDA_SUCCESS + assert torch.equal(torch.ones(1, device="cuda").add_(1).cpu(), torch.full((1,), 2.0)) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py index a526d0752a0a..0da317f4c8ec 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py @@ -31,7 +31,14 @@ from pydantic import ValidationError from transformers import PretrainedConfig -from tensorrt_llm._torch._experimental.modeling_v2._router_index import modeling_v2_resolve +from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( + NoTarget, + engine_facts, + modeling_v2_resolve, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp import ( + routing as qwen4_exp_routing, +) from tensorrt_llm._torch.distributed.ops import MNNVLAllReduce from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_auto import AutoModelForCausalLM @@ -549,13 +556,21 @@ def _qwen4_exp_sm(_on_sm10x): @pytest.mark.usefixtures("_qwen4_exp_sm") @pytest.mark.parametrize( "spec_config", - [None, MTPDecodingConfig(max_draft_len=1), MTPDecodingConfig(max_draft_len=3)], - ids=["no-spec", "mtp1", "mtp3"], + [ + None, + MTPDecodingConfig(max_draft_len=1), + MTPDecodingConfig(max_draft_len=3), + MTPDecodingConfig(max_draft_len=7), + ], + ids=["no-spec", "mtp1", "mtp3", "mtp7"], ) def test_qwen4_exp_tp1_matches(spec_config): config = _model_config(_qwen4_exp_config(), spec_config=spec_config) - target = modeling_v2_resolve(config, _llm_args("require")) - assert _qualified(target).startswith(_QWEN4_EXP_TARGET) + # The recurrent-state cache at either dtype the target serves; `auto` is bf16 here. + for dtype in ("float32", "bfloat16", "auto"): + kv_cache_config = KvCacheConfig(mamba_ssm_cache_dtype=dtype) + target = modeling_v2_resolve(config, _llm_args("require", kv_cache_config=kv_cache_config)) + assert _qualified(target).startswith(_QWEN4_EXP_TARGET), dtype def _mtp(max_draft_len=3, **attributes): @@ -572,14 +587,14 @@ def _mtp(max_draft_len=3, **attributes): "spec_config", [ _NGRAM, - MTPDecodingConfig(max_draft_len=4), + MTPDecodingConfig(max_draft_len=8), MTPDecodingConfig(max_draft_len=3, use_mtp_vanilla=True), _mtp(max_total_draft_tokens=5), _mtp(num_nextn_predict_layers=2), ], ids=[ "ngram", - "draft4", + "draft8", "vanilla-mtp", "max-total-draft", "two-predictors", @@ -606,9 +621,175 @@ def test_qwen4_exp_on_two_ranks_does_not_match(): @pytest.mark.usefixtures("_qwen4_exp_sm") -def test_qwen4_exp_with_an_fp8_kv_cache_does_not_match(): +def test_qwen4_exp_with_an_fp8_ds_mla_kv_cache_does_not_match(): + """fp8_ds_mla is FP8 to the loader but selects the packed sparse-MLA cache layout.""" config = _model_config(_qwen4_exp_config()) - args = dict(kv_cache_config=KvCacheConfig(dtype="fp8")) + args = dict(kv_cache_config=KvCacheConfig(dtype="fp8_ds_mla")) assert modeling_v2_resolve(config, _llm_args("auto", **args)) is None with pytest.raises(ValueError, match="KV cache"): modeling_v2_resolve(config, _llm_args("require", **args)) + + +# --- qwen4_exp: engine choices the constructor would refuse ----------------------- + + +def _declined(config, llm_args_kwargs=None, match=None): + """``auto`` falls back to the built-in model and ``require`` says why.""" + kwargs = llm_args_kwargs or {} + assert modeling_v2_resolve(config, _llm_args("auto", **kwargs)) is None + with pytest.raises(ValueError, match=match): + modeling_v2_resolve(config, _llm_args("require", **kwargs)) + + +def _claimed(config, llm_args_kwargs=None): + target = modeling_v2_resolve(config, _llm_args("require", **(llm_args_kwargs or {}))) + assert _qualified(target).startswith(_QWEN4_EXP_TARGET) + + +@pytest.mark.usefixtures("_qwen4_exp_sm") +@pytest.mark.parametrize("algo", ["fp8", "nvfp4"]) +def test_qwen4_exp_quantized_kv_cache_dtypes_match(algo): + _claimed(_model_config(_qwen4_exp_config()), dict(kv_cache_config=KvCacheConfig(dtype=algo))) + + +@pytest.mark.usefixtures("_qwen4_exp_sm") +@pytest.mark.parametrize("algo", [QuantAlgo.INT8, QuantAlgo.FP8_BLOCK_SCALES]) +def test_qwen4_exp_other_kv_cache_algorithms_are_declined(algo): + facts = engine_facts(_model_config(_qwen4_exp_config()), _llm_args("auto")) + with pytest.raises(NoTarget, match="kv_cache_quant_algo"): + qwen4_exp_routing._check_engine(None, facts._replace(kv_cache_quant_algo=algo)) + + +@pytest.mark.usefixtures("_qwen4_exp_sm") +def test_qwen4_exp_piecewise_prefill_cuda_graph_falls_back(): + config = _model_config(_qwen4_exp_config()) + _declined(config, dict(prefill_cuda_graph_backend="piecewise"), match="piecewise") + + +@pytest.mark.usefixtures("_qwen4_exp_sm") +@pytest.mark.parametrize("backend", ["disabled", "breakable"]) +def test_qwen4_exp_matches_without_or_with_the_breakable_prefill_cuda_graph(backend): + _claimed( + _model_config(_qwen4_exp_config()), + dict(prefill_cuda_graph_backend=backend), + ) + + +@pytest.mark.usefixtures("_qwen4_exp_sm") +@pytest.mark.parametrize("backend", ["CUTLASS", "TRITON", "DEEPGEMM", "VANILLA"]) +def test_qwen4_exp_unimplemented_moe_backend_falls_back(backend): + config = replace(_model_config(_qwen4_exp_config()), moe_backend_requested=backend) + _declined(config, match="moe_config.backend") + + +@pytest.mark.usefixtures("_qwen4_exp_sm") +@pytest.mark.parametrize("backend", ["AUTO", "TRTLLM", "CUTEDSL"]) +def test_qwen4_exp_implemented_moe_backend_matches(backend): + _claimed(replace(_model_config(_qwen4_exp_config()), moe_backend_requested=backend)) + + +@pytest.mark.parametrize( + "spec_kwargs", + [ + dict(max_concurrency=8), + dict(draft_len_schedule={4: 3, 8: 1}), + dict(acceptance_rate_threshold=0.5), + ], + ids=["max-concurrency", "draft-len-schedule", "acceptance-rate"], +) +@pytest.mark.usefixtures("_qwen4_exp_sm") +def test_qwen4_exp_runtime_draft_length_falls_back(spec_kwargs): + """The verify step runs at one width; a draft length that changes at runtime + would trip the width assert in the first reduced step.""" + spec = MTPDecodingConfig(max_draft_len=3, **spec_kwargs) + _declined(_model_config(_qwen4_exp_config(), spec_config=spec), match="draft length") + + +@pytest.mark.usefixtures("_qwen4_exp_sm") +@pytest.mark.parametrize("max_draft_len", [4, 7]) +def test_qwen4_exp_draft_beyond_three_needs_the_cached_replay(monkeypatch, max_draft_len): + """The legacy per-step recurrent states are certified at up to three drafted tokens and + hold one fp32 state per verify row.""" + monkeypatch.setenv("TRTLLM_USE_GDN_REPLAY", "0") + config = _model_config( + _qwen4_exp_config(), spec_config=MTPDecodingConfig(max_draft_len=max_draft_len) + ) + kv_cache_config = KvCacheConfig(mamba_ssm_cache_dtype="float32") + _declined(config, dict(kv_cache_config=kv_cache_config), match="TRTLLM_USE_GDN_REPLAY=0") + monkeypatch.delenv("TRTLLM_USE_GDN_REPLAY") + _claimed(config, dict(kv_cache_config=kv_cache_config)) + + +@pytest.mark.usefixtures("_qwen4_exp_sm") +def test_qwen4_exp_legacy_verify_matches_up_to_three_drafted_tokens(monkeypatch): + monkeypatch.setenv("TRTLLM_USE_GDN_REPLAY", "0") + config = _model_config(_qwen4_exp_config(), spec_config=MTPDecodingConfig(max_draft_len=3)) + _claimed(config, dict(kv_cache_config=KvCacheConfig(mamba_ssm_cache_dtype="float32"))) + + +@pytest.mark.usefixtures("_qwen4_exp_sm") +@pytest.mark.parametrize("max_draft_len", [4, 7]) +def test_qwen4_exp_cutedsl_declines_drafts_beyond_three(monkeypatch, max_draft_len): + """The CuteDSL expert entries capture up to 128 rows (32 requests of 4 verify rows).""" + monkeypatch.delenv("TRTLLM_USE_GDN_REPLAY", raising=False) + kv_cache_config = KvCacheConfig(mamba_ssm_cache_dtype="float32") + spec = MTPDecodingConfig(max_draft_len=max_draft_len) + cutedsl = replace( + _model_config(_qwen4_exp_config(), spec_config=spec), moe_backend_requested="CUTEDSL" + ) + _declined(cutedsl, dict(kv_cache_config=kv_cache_config), match="CUTEDSL") + _claimed( + _model_config(_qwen4_exp_config(), spec_config=spec), dict(kv_cache_config=kv_cache_config) + ) + + +@pytest.mark.usefixtures("_qwen4_exp_sm") +def test_qwen4_exp_cutedsl_matches_up_to_three_drafted_tokens(monkeypatch): + monkeypatch.delenv("TRTLLM_USE_GDN_REPLAY", raising=False) + spec = MTPDecodingConfig(max_draft_len=3) + cutedsl = replace( + _model_config(_qwen4_exp_config(), spec_config=spec), moe_backend_requested="CUTEDSL" + ) + _claimed(cutedsl, dict(kv_cache_config=KvCacheConfig(mamba_ssm_cache_dtype="float32"))) + + +@pytest.mark.usefixtures("_qwen4_exp_sm") +@pytest.mark.parametrize("dtype", ["float32", "bfloat16", "auto"]) +def test_qwen4_exp_mtp_matches_a_float32_or_bfloat16_state_cache(monkeypatch, dtype): + """GDN cached replay runs on either checkpoint dtype; `auto` resolves to bf16 here.""" + monkeypatch.delenv("TRTLLM_USE_GDN_REPLAY", raising=False) + config = _model_config(_qwen4_exp_config(), spec_config=MTPDecodingConfig(max_draft_len=3)) + _claimed(config, dict(kv_cache_config=KvCacheConfig(mamba_ssm_cache_dtype=dtype))) + + +@pytest.mark.parametrize( + "spec_config", [None, MTPDecodingConfig(max_draft_len=3)], ids=["no-spec", "mtp3"] +) +@pytest.mark.usefixtures("_qwen4_exp_sm") +def test_qwen4_exp_declines_a_float16_state_cache(monkeypatch, spec_config): + """The engine's replay gate accepts fp16; the target's pool contract does not.""" + monkeypatch.delenv("TRTLLM_USE_GDN_REPLAY", raising=False) + config = _model_config(_qwen4_exp_config(), spec_config=spec_config) + kv_cache_config = KvCacheConfig(mamba_ssm_cache_dtype="float16") + _declined(config, dict(kv_cache_config=kv_cache_config), match="recurrent-state cache") + + +@pytest.mark.usefixtures("_qwen4_exp_sm") +def test_qwen4_exp_mtp_on_a_float16_state_cache_declines_with_replay_off(monkeypatch): + monkeypatch.setenv("TRTLLM_USE_GDN_REPLAY", "0") + config = _model_config(_qwen4_exp_config(), spec_config=MTPDecodingConfig(max_draft_len=3)) + kv_cache_config = KvCacheConfig(mamba_ssm_cache_dtype="float16") + _declined(config, dict(kv_cache_config=kv_cache_config), match="recurrent-state cache") + + +@pytest.mark.usefixtures("_qwen4_exp_sm") +def test_qwen4_exp_mtp_on_a_bf16_state_cache_matches_with_replay_off(monkeypatch): + monkeypatch.setenv("TRTLLM_USE_GDN_REPLAY", "0") + config = _model_config(_qwen4_exp_config(), spec_config=MTPDecodingConfig(max_draft_len=3)) + _claimed(config) + + +@pytest.mark.usefixtures("_qwen4_exp_sm") +def test_qwen4_exp_no_speculation_matches_the_default_state_cache(monkeypatch): + monkeypatch.delenv("TRTLLM_USE_GDN_REPLAY", raising=False) + _claimed(_model_config(_qwen4_exp_config())) From e10d6c68292ed4b1bb78e3537e5b2e8037331869 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Sun, 18 Oct 2026 22:59:55 -0700 Subject: [PATCH 02/28] [TRTLLM-17033][feat] Qwen3.8-Flash-Next modeling_v2 target: MTP 1..7, CuteDSL experts, video, FP8/NVFP4 KV, bf16 GDN state, breakable prefill graph The target's feature work (`modeling.py`, `weights.py`, its family router `qwen4_exp/routing.py`, the ops it adds or extends) and the unit tests that certify them: - MTP at draft length 1..7 under GDN cached replay (and the legacy verify path up to 4 rows), IndexShare, speculative rollback of the PLE pools and the QSA side caches. - The trunk's routed experts on the CuteDSL sort / grouped-GEMM pipeline (`moe_config.backend: CUTEDSL`). - Video input: tower calls of whole frames, oversize frames refused at admission. - KV storage as the checkpoint's bf16, FP8 or NVFP4 (the QSA paged K/V store and sparse GQA entries for an NVFP4 pool); the Gated-DeltaNet state in fp32 or bf16. - The breakable prefill CUDA graph; canonical QSA selection. - Engine configurations the target cannot run are declined at construction. Folds these commits of the review branch (their changes to the files of this commit): - Run the Qwen3.8-Flash-Next trunk routed experts on CUTEDSL (5405a3d6ee) - Complete MTP for the Qwen3.8-Flash-Next modeling_v2 target (0a1b9c48da) - Accept video input in the Qwen3.8-Flash-Next modeling_v2 target (80cb4ba8b9) - Store the Qwen3.8-Flash-Next KV cache as FP8 when kv_cache_quant_algo is FP8 (744607bf5c) - Decline engine configurations the qwen4_exp target cannot run (311920b003) - Serve video past one tower call and reject an oversize frame at admission (25f5127266) - Cover the qwen4_exp KV storage derivation and attention-call arguments on CPU (d64f96c637) - Check the CuteDSL expert weight path from checkpoint tensors to output (cd749d61d1) - Fold the CUTEDSL expert down_proj input_scale spread into FC2's alpha (158e2afa06) - Tighten the GDN cached-replay op's reference, guard and cell labels (22a43efced) - Let a full IndexShare row give up its oldest ids, not its newest (6d498a45ad) - Hand the prefill target's replay verify the conv rows as packed_qkv (bf56cf0554) - Report K/V magnitudes past the e4m3 range under TRTLLM_MODELING_V2_VALIDATE with an fp8 KV pool (b01202319b) - Remove duplicated and over-long checks from the qwen4_exp target (59a0334e7b) - Run the Qwen3.8-Flash-Next modeling_v2 target at MTP draft length 1..7 under GDN cached replay (bfa50da860) - Serve the Qwen3.8-Flash-Next modeling_v2 Gated-DeltaNet state in bf16 (a0f2206796) - Reword the engine slot-count comments and restate the unit directory's run time (025a17c0ab) - Make the qwen4_exp QSA selection canonical (fc31acbe07) - Serve the qwen4_exp modeling_v2 target under the breakable prefill CUDA graph (0f3c3c078c) - Add the QSA paged K/V store and sparse GQA entries for an NVFP4 KV cache (6ca42e3e6e) - Store the qwen4_exp KV cache as NVFP4 when kv_cache_quant_algo is NVFP4 (e3c62d38f0) - Reword a docstring in the sigmoid-gating branch test (1f449617d9) Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../modeling.py | 3012 ++++++++++++++--- .../ops/activation/gelu_erf.py | 5 +- .../ops/activation/gelu_tanh.py | 8 + .../ops/attention/qsa_canonical_selection.py | 278 ++ .../ops/attention/qsa_expand_block_indices.py | 28 +- .../ops/attention/qsa_mtp_reuse_topk.py | 362 ++ .../ops/attention/qsa_paged_index_scores.py | 73 + .../ops/attention/qsa_paged_index_store.py | 36 +- .../ops/attention/qsa_paged_kv_store.py | 57 +- .../ops/attention/qsa_paged_kv_store_nvfp4.py | 620 ++++ .../ops/attention/qsa_paged_sparse_gqa.py | 43 +- .../attention/qsa_paged_sparse_gqa_nvfp4.py | 1156 +++++++ .../ops/attention/qsa_prefill_compress.py | 36 +- .../ops/attention/qsa_topk_select.py | 32 +- .../ops/attention/vision_rope_2d.py | 15 +- ...sl_nvfp4_gather_grouped_gemm_act_fusion.py | 535 +++ .../ops/moe/cute_dsl_nvfp4_grouped_gemm.py | 457 +++ .../ops/moe/moe_sort.py | 308 ++ .../ops/moe/moe_unpermute_inplace.py | 285 ++ .../ops/norm/layer_norm.py | 6 + .../ops/ple/ngram_token_history.py | 20 +- .../ops/ssm/causal_conv1d_triton_update.py | 44 +- .../ops/ssm/chunk_gated_delta_rule.py | 61 +- .../ops/ssm/cu_seqlens_from_seq_lens.py | 14 +- .../ops/ssm/dilated_short_conv_paged.py | 28 +- ...t_gated_delta_rule_cached_replay_update.py | 668 ++++ ...fused_recurrent_gated_delta_rule_update.py | 56 +- .../fused_sigmoid_gating_delta_rule_update.py | 109 +- .../ops/ssm/reset_recurrent_state_rows.py | 19 +- .../ops/vision/patch_embed_conv3d.py | 16 +- .../weights.py | 123 +- .../modeling_v2/models/qwen4_exp/routing.py | 96 +- ...est_modeling_v2_qsa_canonical_selection.py | 270 ++ ...st_modeling_v2_qsa_expand_block_indices.py | 14 +- .../test_modeling_v2_qsa_mtp_reuse_topk.py | 457 +++ ...test_modeling_v2_qsa_paged_index_scores.py | 21 +- .../test_modeling_v2_qsa_paged_index_store.py | 45 +- .../test_modeling_v2_qsa_paged_kv_store.py | 17 +- ...st_modeling_v2_qsa_paged_kv_store_nvfp4.py | 490 +++ .../test_modeling_v2_qsa_paged_sparse_gqa.py | 84 + ..._modeling_v2_qsa_paged_sparse_gqa_nvfp4.py | 416 +++ .../test_modeling_v2_qsa_prefill_compress.py | 66 +- .../test_modeling_v2_qsa_topk_select.py | 8 +- ...sl_nvfp4_gather_grouped_gemm_act_fusion.py | 711 ++++ ...modeling_v2_cute_dsl_nvfp4_grouped_gemm.py | 733 ++++ .../test_modeling_v2_cute_dsl_weight_path.py | 362 ++ .../moe/test_modeling_v2_moe_sort.py | 432 +++ .../test_modeling_v2_moe_unpermute_inplace.py | 527 +++ ...modeling_v2_causal_conv1d_triton_update.py | 2 +- ...test_modeling_v2_chunk_gated_delta_rule.py | 32 +- ...st_modeling_v2_cu_seqlens_from_seq_lens.py | 10 +- ...st_modeling_v2_dilated_short_conv_paged.py | 15 +- ...t_gated_delta_rule_cached_replay_update.py | 884 +++++ ..._fused_sigmoid_gating_delta_rule_update.py | 170 +- .../test_modeling_v2_gdn_replay_verify.py | 364 ++ .../test_modeling_v2_index_share.py | 202 ++ .../test_modeling_v2_kv_range_probe.py | 90 + .../test_modeling_v2_spec_rollback.py | 435 +++ .../test_modeling_v2_vision_tower_calls.py | 194 ++ .../test_modeling_v2_qwen4_exp_kv_quant.py | 245 ++ .../test_modeling_v2_qwen4_exp_moe_scalars.py | 90 + ...modeling_v2_qwen4_exp_multimodal_embeds.py | 294 ++ ...est_modeling_v2_qwen4_exp_pool_contract.py | 260 +- ...eling_v2_qwen4_exp_thop_attention_calls.py | 131 + ..._modeling_v2_qwen4_exp_validate_request.py | 226 ++ ...test_modeling_v2_qwen4_exp_weight_table.py | 43 +- 66 files changed, 16166 insertions(+), 780 deletions(-) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_canonical_selection.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_mtp_reuse_topk.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store_nvfp4.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_grouped_gemm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_sort.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_unpermute_inplace.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_cached_replay_update.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_canonical_selection.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_mtp_reuse_topk.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store_nvfp4.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_grouped_gemm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_weight_path.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_sort.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_unpermute_inplace.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_gdn_replay_verify.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_index_share.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_kv_range_probe.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_spec_rollback.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_vision_tower_calls.py create mode 100644 tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_kv_quant.py create mode 100644 tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_moe_scalars.py create mode 100644 tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_multimodal_embeds.py create mode 100644 tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_thop_attention_calls.py create mode 100644 tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_validate_request.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py index f71b25463bdd..16944580ba21 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py @@ -11,7 +11,9 @@ when the engine has a ``speculative_config`` (else ``mtp.*`` stays in ``weights.py``'s non-load set). The vision tower is always built; a text-only request never reaches it. Every speculation-dependent change is gated on a host -int that is 1 when speculation is off. +int that is 1 when speculation is off. The PLE pools and the QSA side caches are +written for every verify row, so `SpeculativeRollback` (registered with the spec +worker) rewrites them to the accepted prefix after each verify step. Differences from other targets: @@ -20,10 +22,15 @@ ``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. + prefix; sparse selection replaces the dense core only above it. The MTP draft + layer runs the same branch over its own indexer and its own manager's planes. * **PLE.** Decoder layer 1 only (``ple_layer_ids: [2]`` is one-based): an n-gram side path over a 320,001,536-row fp8 table (47.68 GiB of a 123.5 GiB checkpoint). -* **NVFP4 MoE on every layer** (routed experts only); everything else is bf16. +* **NVFP4 MoE on every layer** (routed experts only); everything else is bf16. The + trunk's routed experts run on trtllm-gen's block-scale runner (``TRTLLM``, and what + ``AUTO`` selects) or on the CuteDSL sort / grouped-GEMM pipeline (``CUTEDSL``), + chosen by ``moe_config.backend``; any other request is refused. The MTP draft + layer's FP8 block-scale experts run on trtllm-gen under either. Integration facts this file depends on: @@ -40,6 +47,37 @@ ``Qwen3VLInputProcessorBase``. This target asserts every ``multimodal_data`` key where it is read. +**KV storage.** The paged K/V pool holds the checkpoint's bf16, fp8-e4m3 when the engine +resolves ``kv_cache_quant_algo: FP8`` (``kv_cache_config.dtype: fp8``), or NVFP4 when it resolves +``NVFP4`` (``kv_cache_config.dtype: nvfp4``): a lossy substitution, since the checkpoint's KV is +bf16 and the fp8 and NVFP4 attention kernels run q, K, V and the softmax probabilities through +e4m3. An NVFP4 pool stores e2m1 pairs (uint8, half a head wide) and keeps one e4m3 scale per 16 +values in a second pool; the second-level scales are 1.0. The NVFP4 attention kernels exist for +the paged-context path only, so an NVFP4 engine needs chunked prefill, block reuse or +speculative decoding. The QSA index and position planes keep their own types under all three. + +**Recurrent-state storage.** The Gated-DeltaNet recurrent pool holds the checkpoint's fp32 +(``mamba_ssm_dtype``) or bf16 when ``kv_cache_config.mamba_ssm_cache_dtype`` asks for it, which +is also what ``auto`` resolves to for this checkpoint (`_SSM_STATE_DTYPES`): a lossy +substitution, since the state is rounded to bf16 once per decode token, once per prefill call and +at every cached-replay checkpoint commit. The conv, PLE and replay-history buffers keep their own +types under both. A bf16 pool makes the decode step call its recurrence without ``cu_seqlens``; +see ``DecodeTarget._gdn_block``. + +**Breakable prefill CUDA graph.** With ``prefill_cuda_graph_backend: breakable`` the engine +captures the core's ``forward`` (the body, not the shell) for a step with context rows at each of +``prefill_capture_num_tokens``, pads any such step up to the nearest bucket and replays it without +calling the body. The shell therefore calls the core's ``begin_step`` before the body on every +step: the asserts, the cache-manager bind, the target choice and the prefill target's host +bookkeeping (``PrefillTarget.prepare``, one immutable ``_PrefillStep``) all live there. The prefill +``forward`` is row-wise over the ``T`` rows it is handed, with four breaks +(``catalog/graph/eager_on_graph``) that run eagerly between the captured segments and read the +step: the vision splice, the PLE block, and each layer's attention core and Gated-DeltaNet +recurrence. Each writes rows ``[:n]`` of a fixed output and zeroes rows ``[n:]``. Under MTP the +prefill body returns one packed ``[T, hidden + stream]`` tensor that ``split_spec_output`` splits, +and the shell hands the spec worker the batch's own ``n`` rows. With the backend off every step is +eager and unpadded and the breaks are plain calls. + **Host offload.** By default the PLE table lives in host memory the device addresses (`weights._host_resident_empty`), and the gather reads it unchanged; ``TRTLLM_MODELING_V2_PLE_HOST_OFFLOAD=0`` restores the resident declaration. Output @@ -64,9 +102,14 @@ fused_qk_norm_rope, ) from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.thop_attention import ( + QM_FP8_KV_CACHE, + QM_NVFP4_KV_CACHE, thop_attention, ) from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm +from tensorrt_llm._torch._experimental.modeling_v2.catalog.graph.eager_on_graph import ( + eager_on_graph, +) from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fp4_block_scale_moe_runner import ( fp4_block_scale_moe_runner, ) @@ -77,26 +120,39 @@ from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_utils import DecoderModel, DecoderModelForCausalLM +from tensorrt_llm._torch.modules.fla.cached_replay import CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE from tensorrt_llm._torch.speculative import get_spec_worker, should_use_separate_draft_kv_cache -from tensorrt_llm._torch.utils import is_gdn_replay_enabled +from tensorrt_llm._torch.utils import ActivationType +from tensorrt_llm.logger import logger +from tensorrt_llm.quantization.mode import QuantAlgo from . import weights as _weights from .ops.activation.fused_sigmoid_mul_inplace import fused_sigmoid_mul_inplace from .ops.activation.gelu_erf import gelu_erf from .ops.activation.gelu_tanh import gelu_tanh +from .ops.attention.qsa_canonical_selection import qsa_canonical_selection from .ops.attention.qsa_expand_block_indices import qsa_expand_block_indices +from .ops.attention.qsa_mtp_reuse_topk import qsa_mtp_reuse_topk from .ops.attention.qsa_paged_index_scores import qsa_paged_index_scores from .ops.attention.qsa_paged_index_store import qsa_paged_index_store from .ops.attention.qsa_paged_kv_store import qsa_paged_kv_store +from .ops.attention.qsa_paged_kv_store_nvfp4 import qsa_paged_kv_store_nvfp4 from .ops.attention.qsa_paged_sparse_gqa import qsa_paged_sparse_gqa +from .ops.attention.qsa_paged_sparse_gqa_nvfp4 import qsa_paged_sparse_gqa_nvfp4 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.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion import ( + cute_dsl_nvfp4_gather_grouped_gemm_act_fusion, +) +from .ops.moe.cute_dsl_nvfp4_grouped_gemm import cute_dsl_nvfp4_grouped_gemm from .ops.moe.fp8_block_scale_moe_runner import fp8_block_scale_moe_runner from .ops.moe.fused_sigmoid_gate_mul_add import fused_sigmoid_gate_mul_add +from .ops.moe.moe_sort import moe_sort +from .ops.moe.moe_unpermute_inplace import moe_unpermute_inplace from .ops.moe.renorm_moe_routing_op import renorm_moe_routing_op from .ops.norm.flashinfer_gemma_rmsnorm import flashinfer_gemma_rmsnorm from .ops.norm.layer_norm import layer_norm @@ -114,6 +170,9 @@ 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_cached_replay_update import ( + fused_recurrent_gated_delta_rule_cached_replay_update, +) from .ops.ssm.fused_recurrent_gated_delta_rule_update import fused_recurrent_gated_delta_rule_update from .ops.ssm.fused_sigmoid_gating_delta_rule_update import fused_sigmoid_gating_delta_rule_update from .ops.ssm.pack_gdn_decode_qkv import pack_gdn_decode_qkv @@ -148,10 +207,10 @@ # Per-call constants of the dense attention call shape: packed-QKV causal GQA, RoPE -# applied outside (fused_qk_norm_rope), bf16 pool. The sparse / skip-softmax group -# stays None/0: QSA calls selection and sparse attention as separate kernels and -# `num_sparse_topk` is 0. `quant_mode` and `predicted_tokens_per_seq` are passed -# beside these. +# applied outside (fused_qk_norm_rope), bf16, fp8-e4m3 or NVFP4 pool. The sparse / skip-softmax +# group stays None/0: QSA calls selection and sparse attention as separate kernels and +# `num_sparse_topk` is 0. `quant_mode`, the two KV scales (all three follow the pool's +# storage) and `predicted_tokens_per_seq` are passed beside these. _CALL_CONSTANTS = dict( output_sf=None, k=None, @@ -163,9 +222,6 @@ mask_type=1, q_scaling=1.0, attention_sinks=None, - # bf16 pool: no KV scale. - kv_scale_orig_quant=None, - kv_scale_quant_orig=None, # In-kernel RoPE disabled; rope_* values are inert placeholders. position_embedding_type=0, rotary_inv_freq=None, @@ -310,6 +366,12 @@ #: Head widths `ops/attention/vision_rope_2d` is certified at. _VISION_ROPE_HEAD_DIMS = (40, 64, 72, 128) +#: The most patch rows one tower call carries, the ceiling the vision ops' cells reach: +#: 8192 placeholder tokens x `spatial_merge_size`**2 patch rows each. `VisionTower` splits +#: a request's frames into calls under it, so the bound a request meets is per frame: one +#: frame's `h * w` patches. +_VISION_CALL_MAX_PATCHES = 32768 + # NVFP4 block size: one e4m3 scale per 16 contiguous elements along K. _SF_VEC = _weights.SF_VEC # The trtllm-gen MoE runner reads activation scales as a *linear* buffer, not the @@ -323,6 +385,18 @@ _ROUTING_METHOD_RENORM = 1 _ACT_TYPE_SWIGLU = 0 +# The trunk's routed-expert backend by the requested `moe_config.backend` +# (`ModelConfig.moe_backend_requested`; `moe_backend` itself is already resolved, +# AUTO -> CUTLASS). Any other request is refused at construction. The MTP draft +# layer's FP8 block-scale experts run on trtllm-gen whatever this says. +_TRUNK_MOE_BACKENDS = {"AUTO": "TRTLLM", "TRTLLM": "TRTLLM", "CUTEDSL": "CUTEDSL"} +# The CuteDSL grouped GEMMs' MMA tile M, which is also the tile `moe_sort` pads each +# expert's rows to: the GEMMs walk the sort's tile tables. sm_103 admits 128 and 256. +_CUTEDSL_TILE_SIZE = 128 +# The longest MTP draft the CuteDSL pipeline is certified at: its cells capture up to 128 rows +# (32 requests of 4 verify rows). +_CUTEDSL_MAX_DRAFT_LEN = 3 + # Skip sentinel the conv ops and paged short conv compare `cache_indices` against # (the runtime's mamba `PAD_SLOT_ID`). Never a substitution: the scheduler's slot # vector gives every row a real slot, and `ops/ple/ngram_token_history` reads its @@ -370,8 +444,9 @@ def _ple_host_offload_enabled() -> bool: # `fused_qk_norm_rope` on `index_scratch` (query side, `use_mrope`) must agree: # both read the three coordinates or neither, else scores mix two position # encodings. Both-on is certified (`mrope_section` (11, 11, 10)); None is a -# modeling choice. No cost on text; on an image batch the index plane sees a frame -# at one position, changing block selection only above `dense_seq_len_threshold`. +# modeling choice. No cost on text; on an image or video batch the index plane sees +# a frame at one position, changing block selection only above +# `dense_seq_len_threshold`. # To flip, set the config's `(s0, s1, s2)` and `_QSA_INDEX_USE_MROPE = True`. _QSA_MROPE_SECTION = None _QSA_INDEX_USE_MROPE = False @@ -390,11 +465,36 @@ def _ple_host_offload_enabled() -> bool: #: FP8 block-scale block edge, on both weight-scale axes and the activation row. _FP8_BLOCK = _weights.FP8_BLOCK +#: Values per e4m3 block scale in an NVFP4 KV pool. +_NVFP4_BLOCK = 16 #: False: plain `amax / 448` scale (the runner's recipe). True: power-of-two #: (UE8M0), a different quantization of the same input. _MTP_MOE_USE_UE8M0 = False +#: The history depth of the GDN cached-replay verify path (the manager's +#: `replay_history_size`). The replay entry is certified at 16 only, and a manager that +#: reports another depth is refused at bind. +_GDN_REPLAY_HISTORY = 16 + +#: The most query rows a verify step carries per generation request, `max_draft_len + 1`. +#: It reaches `thop_attention` as `predicted_tokens_per_seq` and the cached replay as its step +#: width `T`; the replay kernel asserts `T <= 8`, and above it the engine would verify with the +#: legacy per-step states, which this target refuses past `_MAX_LEGACY_VERIFY_WIDTH` rows. +_MAX_VERIFY_WIDTH = 8 + +#: The same bound on the legacy verify path (`TRTLLM_USE_GDN_REPLAY=0`): the entries only it +#: calls (`pack_gdn_decode_qkv`, `fused_gdn_gating`, the multi-token +#: `fused_recurrent_gated_delta_rule_update`) are certified at up to 4 rows per request, and its +#: per-step recurrent buffer holds one state per verify row, request and layer (6.75 GiB at 8 +#: rows and 8 requests in fp32). +_MAX_LEGACY_VERIFY_WIDTH = 4 + +#: The recurrent-state pool dtypes this target serves (`kv_cache_config.mamba_ssm_cache_dtype`): +#: the checkpoint's fp32, or bf16, which rounds the state to bf16 once per decode token, once +#: per prefill call and at every cached-replay checkpoint commit. The engine also accepts fp16. +_SSM_STATE_DTYPES = (torch.float32, torch.bfloat16) + # ----------------------------------------------- KV cache manager binding --- # @@ -402,19 +502,23 @@ def _ple_host_offload_enabled() -> bool: # KV cache manager are `TensorWrapper`s over raw device addresses that do not keep # memory alive, so a view is only valid for the manager that produced it. `_Bound` # is everything derived from one manager; the core replaces it as a unit -# (`_bind_cache_manager`) and the hot path pays one identity comparison. +# (`_bind_cache_manager`) and the hot path pays one identity comparison. The MTP +# draft layer's step can see a second manager (the separate draft KV cache), bound +# the same way into `_DraftBound` (`_bind_draft_cache_manager`). # # The `check_*` functions hold the shape/dtype contract. They read only # `.shape` / `.dtype`, so unit tests run them on meta tensors with no engine. # # For a later target with pools that are not in `attn_metadata` (mamba/GDN state, -# side caches): the reusable parts are (1) one immutable bound object, replaced -# atomically and keyed on manager identity, (2) pure contract checks, (3) refusing -# a re-bind during CUDA graph capture. Once a second target needs it, the next -# step is a declarative `PoolSpec(name, layers, getter, tail_shape, dtype)` list -# consumed by one shared binder under `modeling_v2/` (not `catalog/`), leaving each -# target only its specs and shape constants. Pure-attention targets that read the -# pool pointers from `attn_metadata` on every call (deepseek_v3) do not need this. +# side caches): the reusable parts are (1) one immutable bound object per manager, +# replaced atomically and keyed on manager identity, (2) pure contract checks, +# (3) refusing a bind during CUDA graph capture that no eager step validated first +# (a re-bind for the trunk's, a first bind for the draft layer's). Once a second +# target needs it, the next step is a declarative +# `PoolSpec(name, layers, getter, tail_shape, dtype)` list consumed by one shared +# binder under `modeling_v2/` (not `catalog/`), leaving each target only its specs +# and shape constants. Pure-attention targets that read the pool pointers from +# `attn_metadata` on every call (deepseek_v3) do not need this. class _PoolGeometry(NamedTuple): @@ -427,6 +531,7 @@ class _PoolGeometry(NamedTuple): index_head_dim: int conv_dim: int conv_kernel: int + num_k_heads: int num_v_heads: int head_v_dim: int head_k_dim: int @@ -434,6 +539,29 @@ class _PoolGeometry(NamedTuple): ple_state_len: int ngram_context_len: int dtype: torch.dtype # the activation dtype + kv_dtype: torch.dtype # the paged K/V pool's dtype: bf16, e4m3 (fp8) or uint8 (NVFP4) + kv_width: int # its innermost extent: head_dim, or head_dim // 2 for NVFP4's e2m1 pairs + + +class _KvPlanes(NamedTuple): + """One layer's paged K/V views, and under an NVFP4 pool its e4m3 block-scale views.""" + + k: torch.Tensor + v: torch.Tensor + k_scale: torch.Tensor | None = None + v_scale: torch.Tensor | None = None + + +class _Replay(NamedTuple): + """The GDN cached-replay verify state of one manager: per slot, the recurrent pool + row is a checkpoint and a compact history holds the tokens accepted since it.""" + + old: dict # GDN layer -> (old_x, old_B, old_dt, old_dA_cumsum) history buffers + cache_buf_idx: torch.Tensor # per slot: which history half is pending + prev_num_accepted: torch.Tensor # per slot: pending history entries + history: int + width: int # verify rows per generation request + all_layer_commit: bool # the manager commits partitioned steps over every layer class _Bound(NamedTuple): @@ -443,14 +571,27 @@ class _Bound(NamedTuple): pool_idx: dict # attention layer -> pool offset (what thop_attention addresses) index_cache: dict # attention layer -> QSA index plane position_cache: torch.Tensor # shared QSA position plane - kv: dict # attention layer -> (K, V) views + kv: dict # attention layer -> `_KvPlanes` conv: dict # GDN layer -> conv state pool ssm: dict # GDN layer -> recurrent state pool ple_state: tuple # (PLE conv state, PLE n-gram id history) - # MTP only: the runtime's per-step rollback buffers. + # MTP only: the verify step's per-step conv buffers, and one recurrent mode. Legacy: + # per-step recurrent buffers (`int_ssm`). Cached replay: `replay`, no `int_ssm`. int_idx: torch.Tensor | None = None int_conv: dict | None = None int_ssm: dict | None = None + replay: _Replay | None = None + + +class _DraftBound(NamedTuple): + """The MTP draft layer's views of the manager its draft step sees: the separate + draft manager when the engine built one, else the trunk's.""" + + cm: object + index_cache: torch.Tensor # the draft layer's QSA index plane + position_cache: torch.Tensor # that manager's QSA position plane + kv: _KvPlanes # the draft layer's paged K/V views + share_slot: int | None # its row in the metadata's shared selection buffers (IndexShare) def _require(cm, names, what: str) -> None: @@ -480,15 +621,47 @@ def check_attention_pools(g: _PoolGeometry, layer: int, index_cache, pool) -> No index_shape = (g.tokens_per_block, g.index_kv_heads, g.index_head_dim) _check_pool(index_cache, f"layer {layer}: QSA index plane", 1, index_shape, g.dtype) # The sparse path needs adjacent HND K/V role planes: [pages, 2, H, block, D]. - kv_shape = (g.heads_kv, g.tokens_per_block, g.head_dim) - _check_pool(pool, f"layer {layer}: paged K/V", 2, kv_shape) + kv_shape = (g.heads_kv, g.tokens_per_block, g.kv_width) + _check_pool(pool, f"layer {layer}: paged K/V", 2, kv_shape, g.kv_dtype) assert pool.shape[1] == 2, f"layer {layer}: paged K/V has {pool.shape[1]} role planes, not 2" +def check_kv_scale_planes(g: _PoolGeometry, layer: int, pool, scales) -> None: + """An NVFP4 layer's e4m3 block scales: one uint8 per `_NVFP4_BLOCK` values, laid out like + the data pages (`[slots, 2, H, block, D / 16]`), and holding as many pages per slot as the + data pool does. + + The attention op addresses a layer's data pages and its scale pages with one block-offset + table, so a scale pool whose pages-per-slot differs from the data pool's reads and writes + other slots' scales without any error. Both counts are read off the views' slot strides.""" + scale_shape = (g.heads_kv, g.tokens_per_block, g.head_dim // _NVFP4_BLOCK) + _check_pool(scales, f"layer {layer}: NVFP4 block scales", 2, scale_shape, torch.uint8) + assert scales.shape[1] == 2 and scales.shape[0] == pool.shape[0], ( + f"layer {layer}: the block scales are {tuple(scales.shape)} against paged K/V " + f"{tuple(pool.shape)}; they must hold the same slots, K then V" + ) + pages = [] + for what, t in (("paged K/V", pool), ("block scales", scales)): + page = t[0, 0].numel() + assert t.stride(0) % page == 0, ( + f"layer {layer}: the {what} slot stride {t.stride(0)} is not a whole number of " + f"{page}-element pages" + ) + pages.append(t.stride(0) // page) + assert pages[0] == pages[1], ( + f"layer {layer}: the paged K/V pool holds {pages[0]} pages per slot and the block-scale " + f"pool {pages[1]}; the attention op would address the scales of other slots" + ) + + def check_recurrent_pools(g: _PoolGeometry, conv, ssm) -> None: _check_pool(conv, "conv state pool", 1, (g.conv_dim, g.conv_kernel - 1), g.dtype) # Key dim innermost. _check_pool(ssm, "recurrent state pool", 1, (g.num_v_heads, g.head_v_dim, g.head_k_dim)) + assert ssm.dtype in _SSM_STATE_DTYPES, ( + f"the recurrent state pool is {ssm.dtype}; this target serves " + f"{' or '.join(str(d) for d in _SSM_STATE_DTYPES)} (kv_cache_config.mamba_ssm_cache_dtype)" + ) def check_ple_pools(g: _PoolGeometry, ple_conv, ple_ngram, conv) -> None: @@ -501,22 +674,165 @@ def check_ple_pools(g: _PoolGeometry, ple_conv, ple_ngram, conv) -> None: ) -def check_step_buffers(g: _PoolGeometry, int_conv, int_ssm, conv, ssm, int_idx) -> None: - """MTP verify rollback buffers: `[slots, steps, ...]`, at the pools' dtypes - (the promotion kernel copies between them without converting).""" - assert int_conv is not None and int_ssm is not None, ( - "the runtime allocated no per-step conv/recurrent state buffers; a " - "partially accepted draft cannot be rolled back" +def check_step_buffers(g: _PoolGeometry, int_conv, conv, int_idx) -> None: + """MTP verify buffers both recurrent modes use: the per-step conv windows + `[slots, steps, ...]` at the pool's dtype (the promotion copies between them without + converting) and the rows `update_mamba_states` reads its sources from.""" + assert int_conv is not None, ( + "the runtime allocated no per-step conv buffer; a partially accepted draft " + "cannot be rolled back" ) _check_pool(int_conv, "per-step conv buffer", 2, (g.conv_dim, g.conv_kernel - 1), conv.dtype) - ssm_shape = (g.num_v_heads, g.head_v_dim, g.head_k_dim) - _check_pool(int_ssm, "per-step recurrent buffer", 2, ssm_shape, ssm.dtype) - # `update_mamba_states` reads its source rows from this vector. assert int_idx is not None and int_idx.dtype == torch.int32, ( "the cache manager exposes no int32 intermediate_state_indices" ) +def check_legacy_state(g: _PoolGeometry, int_ssm, ssm) -> None: + """Legacy verify: per-step recurrent buffer `[slots, steps, ...]` at the pool's dtype.""" + assert int_ssm is not None, ( + "the cache manager runs the legacy verify path but allocated no per-step recurrent " + "buffer; a partially accepted draft cannot be rolled back" + ) + ssm_shape = (g.num_v_heads, g.head_v_dim, g.head_k_dim) + _check_pool(int_ssm, "per-step recurrent buffer", 2, ssm_shape, ssm.dtype) + assert int_ssm.shape[1] <= _MAX_LEGACY_VERIFY_WIDTH, ( + f"the legacy verify path runs {int_ssm.shape[1]} rows per generation request; its entries " + f"are certified at up to {_MAX_LEGACY_VERIFY_WIDTH} (max_draft_len " + f"{_MAX_LEGACY_VERIFY_WIDTH - 1}). Leave TRTLLM_USE_GDN_REPLAY unset (cached replay) for " + "a longer draft" + ) + + +def check_replay_state(history, width, prev_accepted, cache_buf_idx, ssm) -> None: + """Cached replay, manager-wide: the geometry the replay entry is certified at, and the + per-slot bookkeeping, which must cover every pool slot.""" + assert ssm.dtype in _SSM_STATE_DTYPES, ( + f"the recurrent pool is {ssm.dtype}; the replay entry is certified at an fp32 or a " + "bf16 checkpoint" + ) + assert history == _GDN_REPLAY_HISTORY, ( + f"replay_history_size {history}; the replay entry is certified at {_GDN_REPLAY_HISTORY}" + ) + assert 2 <= width <= min(_MAX_VERIFY_WIDTH, history), ( + f"replay_step_width {width}; the replay entry is certified at 2..{_MAX_VERIFY_WIDTH} " + f"verify rows and never above the history ({history})" + ) + for name, t in (("prev_num_accepted_tokens", prev_accepted), ("cache_buf_idx", cache_buf_idx)): + assert t.dim() == 1 and t.shape[0] >= ssm.shape[0] and t.dtype == torch.int32, ( + f"{name} is {tuple(t.shape)} {t.dtype}; the kernel reads one int32 per pool slot " + f"({ssm.shape[0]})" + ) + + +def check_replay_history( + g: _PoolGeometry, layer: int, history, pool_slots, old_x, old_B, old_dt +) -> None: + """One GDN layer's replay history: `[slots, 2 halves, ...]`, addressed by `history` and + the head geometry, never by the buffer's own shape.""" + for name, t in (("old_x", old_x), ("old_B", old_B), ("old_dt", old_dt)): + assert t is not None, f"layer {layer}: the manager's layer cache carries no {name}" + assert old_x.shape[0] >= pool_slots, ( + f"layer {layer}: the replay history has {old_x.shape[0]} slots for a {pool_slots}-slot pool" + ) + bf16, fp32 = torch.bfloat16, torch.float32 + x_tail, b_tail = (history, g.num_v_heads, g.head_v_dim), (history, g.num_k_heads, g.head_k_dim) + _check_pool(old_x, f"layer {layer}: old_x", 2, x_tail, bf16) + _check_pool(old_B, f"layer {layer}: old_B", 2, b_tail, bf16) + _check_pool(old_dt, f"layer {layer}: old_dt", 2, (g.num_v_heads, history), fp32) + assert old_x.shape[1] == old_B.shape[1] == old_dt.shape[1] == 2, ( + f"layer {layer}: the replay history is not double-buffered" + ) + + +def check_verify_width(rows: int) -> None: + """The rows per generation request one step carries: `max_draft_len + 1` (`runtime_draft_len + + 1` at run time), a `thop_attention` `predicted_tokens_per_seq` the entry is certified at.""" + assert 1 <= rows <= _MAX_VERIFY_WIDTH, ( + f"predicted_tokens_per_seq {rows} is outside the 1..{_MAX_VERIFY_WIDTH} range " + "attention/thop_attention is certified at" + ) + + +def check_replay_width(replay, rows: int) -> None: + """Under the cached replay a single-row generation step cannot take the ordinary + recurrence: the pool row is a checkpoint whose pending history only the replay kernel + applies. The manager's step width is the one row count per sequence a step may carry.""" + assert replay is None or rows == replay.width, ( + f"the batch carries {rows} query rows per generation sequence and the GDN cached " + f"replay runs at step width {replay.width}; a runtime_draft_len of {rows - 1} (a " + "draft-length schedule, or skipped drafting) cannot run under it. Export " + "TRTLLM_USE_GDN_REPLAY=0 for variable draft lengths" + ) + + +def check_shared_selection(shared, slots, layer: int, block_topk: int) -> None: + """IndexShare: the metadata's shared selection buffers have a row for `layer` and are + `block_topk` wide.""" + assert shared is not None, "IndexShare is on but the QSA metadata allocated no shared buffers" + assert layer in slots and slots[layer] < shared.shape[0], ( + f"the KV cache manager gives layer {layer} no row of the {shared.shape[0]}-row shared " + f"selection buffers (qsa_shared_topk_slots={slots})" + ) + assert shared.shape[2] == block_topk, ( + f"the shared selection rows are {shared.shape[2]} wide and this target selects " + f"{block_topk} blocks per row" + ) + + +def _moe_layer_scalars(w, i: int, backend: str) -> dict[str, torch.Tensor]: + """Layer `i`'s routed-expert scalars for `backend`, built from the checkpoint's + `input_scale` / `weight_scale_2` rows. + + The checkpoint stores modelopt reciprocals (`input_scale = 1 / g_act`, + `weight_scale_2 = 1 / g_w`), so a quantizer's global scale is `1 / input_scale` and + the runner's scalars are products of the two. + """ + isc1, isc1_up = w[f"l{i}_e_isc1"], w[f"l{i}_e_isc1_up"] + assert torch.equal(isc1, isc1_up), ( + f"layer {i}: expert gate/up input_scale differ; the fused FC1 " + "reads one activation quantization for both halves" + ) + assert torch.equal(w[f"l{i}_e_ws2_1"], w[f"l{i}_e_ws2_1_up"]), ( + f"layer {i}: expert gate/up weight_scale_2 differ; the fused " + "FC1 carries one alpha per expert for both halves" + ) + # One activation quantization feeds all 512 experts, so FC1's activation + # global scale must be one number (the shared expert is bf16, no scale). + assert torch.equal(isc1.min(), isc1.max()), ( + f"layer {i}: the routed experts do not share one " + "gate_proj.input_scale, so a single activation quantization " + "cannot feed them all and this assembly is wrong for this " + "checkpoint" + ) + e_isc2, ws2_2 = w[f"l{i}_e_isc2"], w[f"l{i}_e_ws2_2"] + gate1 = (isc1 * w[f"l{i}_e_ws2_1"]).contiguous() + # Routed FC1 activation global scale (one per layer). + scalars = dict(global_scale=(1.0 / isc1[:1]).contiguous()) + if backend == "CUTEDSL": + # The CuteDSL FC1 epilogue requantizes with one global scale for the whole layer, + # where trtllm-gen folds each expert's `down_proj.input_scale` into its FC1 + # scalar. Quantizing at `1 / max(input_scale)` (the smallest per-expert global + # scale, so no block scale saturates) and putting the same factor into FC2's alpha + # makes `global_sf * alpha2[e] == ws2_2[e]` for every expert, whatever the spread. + isc2_max = e_isc2.max() + scalars.update( + # [num_experts] each: activation scale times weight scale. + alpha1=gate1, + alpha2=(isc2_max * ws2_2).contiguous(), + # The FC2 activation's quantization scale, numel 1. + global_sf=(1.0 / isc2_max).reshape(1).contiguous(), + ) + else: + scalars.update( + # [num_experts] each. Swapping the first two is finite and silent. + output1_scale_scalar=(gate1 / e_isc2).contiguous(), + output1_scale_gate_scalar=gate1, + output2_scale_scalar=(e_isc2 * ws2_2).contiguous(), + ) + return scalars + + class Qwen4ExpTp1ModelingV2(DecoderModel): def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig): """`composite_config` is passed in because the text sub-config adopted as @@ -769,6 +1085,10 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig f"seq_len_threshold above token_topk is the way to get here, and " f"it is a configuration this target does not implement" ) + # IndexShare for the MTP draft loop (`sparse_attention_config. + # index_share_for_mtp_iteration`), as the engine resolved it: the same lowering + # sizes the metadata's shared selection buffers. + self.mtp_index_share = bool(qsa_params.mtp_index_share) self.index_width = (self.index_heads + self.index_kv_heads) * self.index_head_dim # Packed [q | k | v] scratch for the fused norm+rope call; the spare KV head # is the v slot the op needs but never reads. @@ -810,7 +1130,9 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig # Every layer is MoE; the dense `intermediate_size` is dead config. assert getattr(cfg, "decoder_sparse_step", 1) == 1, "every layer must be MoE" assert not getattr(cfg, "mlp_only_layers", []), "every layer must be MoE" - # trtllm-gen NVFP4 runner operand-width rules and fp4_quantize's block rule. + # trtllm-gen NVFP4 runner operand-width rules and fp4_quantize's block rule; they + # also cover the CuteDSL FC1 / FC2 rules (K and I in 64-element scale groups, the + # gated N in 128-row pairs of 64 up rows and 64 gate rows). assert self.hidden % 256 == 0, "fp4_block_scale_moe_runner hidden alignment" assert (2 * self.moe_inter) % 128 == 0, "fp4_block_scale_moe_runner 2I alignment" assert (self.hidden // _SF_VEC) % 4 == 0 and (self.moe_inter // _SF_VEC) % 4 == 0, ( @@ -818,6 +1140,14 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig ) # flashinfer_silu_and_mul's precondition is on the half width. assert (self.shared_inter * 2) % 16 == 0, "flashinfer_silu_and_mul half alignment" + # Which kernels run the trunk's routed experts; it also fixes the load-time layout + # of the expert stacks (`weights.py`), so it is decided before anything is declared. + requested = model_config.moe_backend_requested.upper() + assert requested in _TRUNK_MOE_BACKENDS, ( + f"moe_config.backend={model_config.moe_backend_requested!r} is not implemented by " + f"this target; the supported requests are {sorted(_TRUNK_MOE_BACKENDS)}" + ) + self.trunk_moe_backend = _TRUNK_MOE_BACKENDS[requested] # ---- PLE ----------------------------------------------------------- ple_ids = list(getattr(cfg, "ple_layer_ids", None) or []) @@ -931,16 +1261,32 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig "must publish them as `multimodal_token_ids`" ) - # ---- KV pool dtype -------------------------------------------------- - # An unsupported KV dtype makes the QSA backend silently fall back to dense - # (`QSA_SPARSE_KV_CACHE_DTYPES`). + # ---- KV pool storage ------------------------------------------------ + # `None` keeps the checkpoint's bf16 K/V (`quant_mode` 0); `FP8`, what + # `kv_cache_config.dtype: fp8` resolves to, stores them as e4m3 under + # QM_FP8_KV_CACHE; `NVFP4` as e2m1 pairs with e4m3 block scales under + # QM_NVFP4_KV_CACHE. Anything else is refused: the QSA backend serves its sparse path + # over half, bf16, fp8 and (with its own kernels here) NVFP4 pools, and silently falls + # back to dense otherwise. The index and position planes are kernel contracts, not KV + # storage, and keep their types under all three. kv_algo = model_config.quant_config.kv_cache_quant_algo - assert kv_algo is None, ( - f"this target runs the checkpoint's own bf16 KV pool, which is what " - f"the accuracy reference was measured at and what the QSA sparse " - f"path requires; the engine resolved kv_cache_quant_algo {kv_algo!r}" - ) - self.quant_mode = 0 + assert kv_algo in (None, QuantAlgo.FP8, QuantAlgo.NVFP4), ( + f"this target stores K/V as the checkpoint's bf16 (kv_cache_quant_algo None), " + f"as fp8-e4m3 (FP8) or as NVFP4; the engine resolved kv_cache_quant_algo {kv_algo!r}" + ) + self.quant_mode = { + None: 0, + QuantAlgo.FP8: QM_FP8_KV_CACHE, + QuantAlgo.NVFP4: QM_NVFP4_KV_CACHE, + }[kv_algo] + # What `check_attention_pools` holds every manager's K/V pool to, the draft layer's + # included: the dtype and the innermost extent (an NVFP4 byte holds two values). + self.kv_pool_dtype = { + None: dt, + QuantAlgo.FP8: torch.float8_e4m3fn, + QuantAlgo.NVFP4: torch.uint8, + }[kv_algo] + self.kv_pool_width = self.head_dim // 2 if kv_algo == QuantAlgo.NVFP4 else self.head_dim # The weight table reads the tower's geometry off this. self.vision_config = vcfg @@ -961,14 +1307,18 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig f"the checkpoint; this target carries the weights of " f"{_MTP_NUM_LAYERS} MTP module" ) - # `max_draft_len` becomes `predicted_tokens_per_seq` on every attention - # call, certified at 1..4, so attention binds. + # `max_draft_len + 1` becomes `predicted_tokens_per_seq` on every attention call + # and the cached replay's step width, certified at up to `_MAX_VERIFY_WIDTH`. max_draft = spec_config.max_draft_len - assert 1 <= max_draft <= 3, ( - f"max_draft_len={max_draft} needs predicted_tokens_per_seq " - f"{max_draft + 1}, and attention/thop_attention certifies that " - f"argument at 1..4. A longer draft needs a kernel-verify " - f"dispatch, not a bigger number" + assert 1 <= max_draft <= _MAX_VERIFY_WIDTH - 1, ( + f"max_draft_len={max_draft} needs {max_draft + 1} verify rows per request, and " + f"attention/thop_attention and the GDN cached replay are certified at up to " + f"{_MAX_VERIFY_WIDTH} (max_draft_len {_MAX_VERIFY_WIDTH - 1}); 8 is also the " + "replay kernel's ceiling" + ) + assert self.trunk_moe_backend != "CUTEDSL" or max_draft <= _CUTEDSL_MAX_DRAFT_LEN, ( + f"max_draft_len={max_draft} with moe_config.backend=CUTEDSL: the CuteDSL expert " + f"pipeline is certified at max_draft_len up to {_CUTEDSL_MAX_DRAFT_LEN}" ) # A linear tree keeps the Blackwell spec-dec mask surface inert. assert not getattr(spec_config, "use_dynamic_tree", False), ( @@ -983,19 +1333,10 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig f"{spec_config.max_total_draft_tokens} against {max_draft}" ) - # The engine picks the replay (default) or legacy verify path at - # cache-manager construction; this target is closed over the legacy one - # (`_gdn_verify` + `update_mamba_states`). Checked again against the - # manager at the first forward. - assert not is_gdn_replay_enabled(), ( - "the GDN cached-replay verify path is on and this target does " - "not implement it: its commit kernel " - "(fused_recurrent_gated_delta_rule_cached_replay_update) has no " - "catalog entry, and with replay on the runtime allocates no " - "intermediate SSM buffer for the legacy path to write. Export " - "TRTLLM_USE_GDN_REPLAY=0 before the ranks start. This is a " - "requirement of the MTP path, not a tuning knob" - ) + # The recurrent half of the verify step has two forms, and the cache manager + # picks one when it is built (`TRTLLM_USE_GDN_REPLAY`, default on): the cached + # replay or the legacy per-step recurrent buffers. Both are implemented; + # `_bind_step_buffers` reads which one a manager allocated. # The FP8 block-scale runner's own shape rules. assert self.hidden % _FP8_BLOCK == 0 and self.moe_inter % _FP8_BLOCK == 0, ( f"fp8_block_scale_moe_runner needs hidden ({self.hidden}) and " @@ -1023,8 +1364,24 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig self._moe_scalars: list | None = None self._gdn_gating: dict | None = None self._cos_sin: torch.Tensor | None = None + # The attention op's KV scales: fp32 [1] ones over an fp8 pool (never None there: the + # generation kernels read `kvCacheScale[0]`), fp32 [3] ones (q, K, V) over an NVFP4 + # one, None over bf16. Unit, because the fp8 QSA store and sparse GQA take no scale + # while a request's history is written by the dense and the sparse path alike, so the + # two agree on what a slot holds only at 1.0 (the NVFP4 entries take the [3] tensor + # and agree at any value, but the checkpoint has no KV scales to calibrate with). + # Lossy: fp8 K/V are stored by a plain cast, so a magnitude past the e4m3 maximum (448) + # saturates (`_KvRangeProbe` reports it under TRTLLM_MODELING_V2_VALIDATE). + self.kv_scale_orig_quant: torch.Tensor | None = None + self.kv_scale_quant_orig: torch.Tensor | None = None # MTP module operands; None with speculation off. self._mtp: _MtpWeights | None = None + self._mtp_cos_sin: torch.Tensor | None = None + # The verify step's side-cache rollback, registered with the spec worker; None + # with speculation off. + self.spec_rollback: SpeculativeRollback | None = ( + SpeculativeRollback() if self.mtp_enabled else None + ) self._prefill: "PrefillTarget | None" = None self._decode: "DecodeTarget | None" = None self.vision_tower: VisionTower | None = None @@ -1032,6 +1389,10 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig # None until the first forward. self._bound: _Bound | None = None self._binds = 0 + # Set by `begin_step`, cleared by `forward`: a forward without a prepared step is a bug. + self._step_begun = False + # The draft layer's own bind, keyed on the manager its draft step sees. + self._draft_bound: _DraftBound | None = None # ---------------------------------------------------------------- load --- @@ -1045,40 +1406,9 @@ def post_load_weights(self) -> None: Not an override (DecoderModel has no such hook); named for the shell hook.""" w = self.w device = w["embed"].device - moe_scalars = [] - for i in range(self.num_layers): - # The checkpoint stores modelopt reciprocals (`input_scale = 1 / g_act`, - # `weight_scale_2 = 1 / g_w`), so the quantizer's global scale is - # `1 / input_scale` and the runner's scalars are products of the two. - isc1, isc1_up = w[f"l{i}_e_isc1"], w[f"l{i}_e_isc1_up"] - assert torch.equal(isc1, isc1_up), ( - f"layer {i}: expert gate/up input_scale differ; the fused FC1 " - "reads one activation quantization for both halves" - ) - assert torch.equal(w[f"l{i}_e_ws2_1"], w[f"l{i}_e_ws2_1_up"]), ( - f"layer {i}: expert gate/up weight_scale_2 differ; the fused " - "FC1 carries one alpha per expert for both halves" - ) - # One activation quantization feeds all 512 experts, so FC1's activation - # global scale must be one number (the shared expert is bf16, no scale). - assert torch.equal(isc1.min(), isc1.max()), ( - f"layer {i}: the routed experts do not share one " - "gate_proj.input_scale, so a single activation quantization " - "cannot feed them all and this assembly is wrong for this " - "checkpoint" - ) - e_isc2 = w[f"l{i}_e_isc2"] - gate1 = (isc1 * w[f"l{i}_e_ws2_1"]).contiguous() - moe_scalars.append( - dict( - # [num_experts] each. Swapping the first two is finite and silent. - output1_scale_scalar=(gate1 / e_isc2).contiguous(), - output1_scale_gate_scalar=gate1, - output2_scale_scalar=(e_isc2 * w[f"l{i}_e_ws2_2"]).contiguous(), - # Routed FC1 activation global scale (one per layer). - global_scale=(1.0 / isc1[:1]).contiguous(), - ) - ) + moe_scalars = [ + _moe_layer_scalars(w, i, self.trunk_moe_backend) for i in range(self.num_layers) + ] 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. @@ -1107,15 +1437,20 @@ def post_load_weights(self) -> None: f"do not concatenate to the padded id space this checkpoint " f"declares" ) - # The index branch's rope table in the QSA kernels' layout: - # `[max_positions, 2 * rotary_pairs]` fp32, cos block then sin block. - pairs = self.rotary_dim // 2 - exponent = torch.arange(0, self.rotary_dim, 2, dtype=torch.float64, device=device) - inv_freq = (1.0 / (self.theta ** (exponent / self.rotary_dim))).to(torch.float32) - positions = torch.arange(self.max_positions, dtype=torch.float32, device=device) - angles = positions.unsqueeze(1) * inv_freq.unsqueeze(0) - self._cos_sin = torch.cat([torch.cos(angles), torch.sin(angles)], dim=1).contiguous() - assert self._cos_sin.shape == (self.max_positions, 2 * pairs), self._cos_sin.shape + self._cos_sin = self._index_rope_table(self.theta, device) + if self.mtp_enabled: + # The draft layer's index plane rotates at the module's own rope base, like its + # attention rope. + self._mtp_cos_sin = ( + self._cos_sin + if self.mtp_theta == self.theta + else self._index_rope_table(self.mtp_theta, device) + ) + if self.quant_mode in (QM_FP8_KV_CACHE, QM_NVFP4_KV_CACHE): + # Built here, not at construction: that runs under meta init. + scales = 3 if self.quant_mode == QM_NVFP4_KV_CACHE else 1 + self.kv_scale_quant_orig = torch.ones(scales, dtype=torch.float32, device=device) + self.kv_scale_orig_quant = 1.0 / self.kv_scale_quant_orig self._report_ple_residency() if self.mtp_enabled: self._mtp = _mtp_views(self) @@ -1124,6 +1459,18 @@ def post_load_weights(self) -> None: self._decode = DecodeTarget(self) self.vision_tower = VisionTower(self) + def _index_rope_table(self, theta: float, device) -> torch.Tensor: + """The index branch's rope table at `theta` in the QSA kernels' layout: + `[max_positions, 2 * rotary_pairs]` fp32, cos block then sin block.""" + pairs = self.rotary_dim // 2 + exponent = torch.arange(0, self.rotary_dim, 2, dtype=torch.float64, device=device) + inv_freq = (1.0 / (theta ** (exponent / self.rotary_dim))).to(torch.float32) + positions = torch.arange(self.max_positions, dtype=torch.float32, device=device) + angles = positions.unsqueeze(1) * inv_freq.unsqueeze(0) + table = torch.cat([torch.cos(angles), torch.sin(angles)], dim=1).contiguous() + assert table.shape == (self.max_positions, 2 * pairs), table.shape + return table + def _report_ple_residency(self) -> None: """Check where the n-gram table ended up, and report it once. @@ -1172,6 +1519,7 @@ def _pool_geometry(self, tokens_per_block: int) -> _PoolGeometry: index_head_dim=self.index_head_dim, conv_dim=self.conv_dim, conv_kernel=self.conv_kernel, + num_k_heads=self.num_k_heads, num_v_heads=self.num_v_heads, head_v_dim=self.head_v_dim, head_k_dim=self.head_k_dim, @@ -1179,7 +1527,21 @@ def _pool_geometry(self, tokens_per_block: int) -> _PoolGeometry: ple_state_len=self.ple_state_len, ngram_context_len=self.ngram_context_len, dtype=self.w["embed"].dtype, - ) + kv_dtype=self.kv_pool_dtype, + kv_width=self.kv_pool_width, + ) + + def _kv_planes(self, cm, g: _PoolGeometry, layer: int, index_cache) -> _KvPlanes: + """One layer's paged K/V views from `cm`, checked against the contract: the data + planes, and under NVFP4 the block-scale planes that go with them.""" + pool = cm.get_buffers(layer, kv_layout="HND") + check_attention_pools(g, layer, index_cache, pool) + if self.quant_mode != QM_NVFP4_KV_CACHE: + return _KvPlanes(pool[:, 0], pool[:, 1]) + _require(cm, ("get_block_scale_buffers",), "reads an NVFP4 pool's e4m3 block-scale planes") + scales = cm.get_block_scale_buffers(layer, kv_layout="HND") + check_kv_scale_planes(g, layer, pool, scales) + return _KvPlanes(pool[:, 0], pool[:, 1], scales[:, 0], scales[:, 1]) def _bind_cache_manager(self, md, position_ids) -> None: """Bind this module to one KV cache manager: fail-fast on the contract, then @@ -1239,6 +1601,13 @@ def _bind_cache_manager(self, md, position_ids) -> None: "a sliding window is not implemented here" ) + # The NVFP4 attention kernels exist for the paged-context path only; without it the + # dense op would not read the pool the sparse path wrote. + assert self.quant_mode != QM_NVFP4_KV_CACHE or md.use_paged_context_fmha, ( + "an NVFP4 KV cache is read by the paged-context attention kernels only, which the " + "engine enables for chunked prefill, KV cache block reuse or speculative decoding; " + "none of the three is on" + ) g = self._pool_geometry(md.tokens_per_block) _require( cm, @@ -1257,9 +1626,7 @@ def _bind_cache_manager(self, md, position_ids) -> None: index_cache, kv = {}, {} for i in self.attn_layers: index_cache[i] = cm.get_index_k_buffer(i) - pool = cm.get_buffers(i, kv_layout="HND") - check_attention_pools(g, i, index_cache[i], pool) - kv[i] = (pool[:, 0], pool[:, 1]) + kv[i] = self._kv_planes(cm, g, i, index_cache[i]) # A missed hybrid predicate would still build and pass a smoke run with 36 # layers carrying no state. @@ -1305,33 +1672,113 @@ def _bind_cache_manager(self, md, position_ids) -> None: cm, pool_idx, index_cache, position_cache, kv, conv, ssm, (ple_conv, ple_ngram) ) if self.mtp_enabled: - bound = bound._replace(**self._bind_step_buffers(cm, g, conv0, ssm0)) + bound = bound._replace(**self._bind_step_buffers(cm, g, mm, conv0, ssm0)) self._bound = bound - def _bind_step_buffers(self, cm, g: _PoolGeometry, conv0, ssm0) -> dict: + def _bind_step_buffers(self, cm, g: _PoolGeometry, mm, conv0, ssm0) -> dict: """Per-manager fail-fast for the MTP increment: the verify step needs - runtime-owned per-step buffers that exist only under the legacy GDN path. - Returns the `_Bound` fields it adds.""" + runtime-owned buffers the base path never touches. Every row's conv window goes + to a per-step buffer in both recurrent modes; the recurrent state is either the + cached replay (a checkpoint per slot plus a compact history of accepted tokens, + no per-step recurrent buffer) or, with `TRTLLM_USE_GDN_REPLAY=0`, per-step + recurrent buffers. Returns the `_Bound` fields it adds.""" assert cm.is_speculative(), ( f"the KV cache manager {type(cm).__name__} reports is_speculative() False, " "so it allocated none of the per-step state buffers the verify path writes" ) - assert not getattr(cm, "use_replay_state_update", False), ( - "the GDN cached-replay verify path is on. Export TRTLLM_USE_GDN_REPLAY=0 " - "before the ranks start; see this target's __init__ for why" - ) _require( cm, ("get_intermediate_conv_states", "get_intermediate_ssm_states"), - "reads both per-step state buffers", + "binds its per-step state buffers", ) # `get_intermediate_*` builds a fresh view per call, so bind them here. int_conv = {i: cm.get_intermediate_conv_states(i) for i in self.gdn_layers} - int_ssm = {i: cm.get_intermediate_ssm_states(i) for i in self.gdn_layers} int_idx = getattr(cm, "intermediate_state_indices", None) g0 = self.gdn_layers[0] - check_step_buffers(g, int_conv[g0], int_ssm[g0], conv0, ssm0, int_idx) - return dict(int_idx=int_idx, int_conv=int_conv, int_ssm=int_ssm) + check_step_buffers(g, int_conv[g0], conv0, int_idx) + int_ssm0 = cm.get_intermediate_ssm_states(g0) + if not getattr(cm, "use_replay_state_update", False): + check_legacy_state(g, int_ssm0, ssm0) + int_ssm = {i: cm.get_intermediate_ssm_states(i) for i in self.gdn_layers} + return dict(int_idx=int_idx, int_conv=int_conv, int_ssm=int_ssm) + assert int_ssm0 is None, ( + "the cache manager runs the cached-replay verify path and also allocated a " + "per-step recurrent buffer; the runtime allocates one or the other" + ) + return dict(int_idx=int_idx, int_conv=int_conv, replay=self._bind_replay(cm, g, mm, ssm0)) + + def _bind_replay(self, cm, g: _PoolGeometry, mm, ssm0) -> _Replay: + """The cached-replay history buffers and bookkeeping of one manager.""" + meta = cm.get_replay_state_update_metadata() + assert meta is not None, ( + "use_replay_state_update is True but the manager returns no replay metadata " + "(prev_num_accepted_tokens / cache_buf_idx / widths)" + ) + history, width = meta.replay_history_size, meta.replay_step_width + check_replay_state(history, width, meta.prev_num_accepted_tokens, meta.cache_buf_idx, ssm0) + old = {} + for i in self.gdn_layers: + layer_cache = cm.mamba_layer_cache(i) + old[i] = ( + layer_cache.old_x, + layer_cache.old_B, + layer_cache.old_dt, + layer_cache.old_dA_cumsum, + ) + check_replay_history(g, i, history, ssm0.shape[0], *old[i][:3]) + # The partitioned replay is the manager's to commit: it commits exactly when this + # flag holds and a step carries `CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE` or more + # generation requests, and the forward hands the kernel the work items under that + # same condition. + all_layer = bool(getattr(cm, "use_gdn_cached_replay_all_layer_commit", False)) + if all_layer: + missing = [n for n in ("replay_work_items", "replay_n_writes") if not hasattr(mm, n)] + assert not missing, ( + f"the manager's all-layer replay commit is on but the mamba metadata lacks " + f"{missing}; the partitioned verify call needs both" + ) + return _Replay( + old, meta.cache_buf_idx, meta.prev_num_accepted_tokens, history, width, all_layer + ) + + def _bind_draft_cache_manager(self, md) -> None: + """Bind the draft layer to the manager its draft step sees: its own index plane, + that manager's position plane and its paged K/V. + + With a separate draft KV cache that is the draft manager the worker swaps in for + the draft loop, whose pools hold the draft layer's pages and none of the trunk's; + otherwise the trunk's manager, which maps the draft layer as well. Keyed on the + manager like `_bind_cache_manager`. Never first bound under capture: the runner + runs every captured batch shape eagerly first, so a bind under capture would bake + views no eager step validated into the graph.""" + cm = md.kv_cache_manager + assert not torch.cuda.is_current_stream_capturing(), ( + f"the draft layer is not bound to KV cache manager {type(cm).__name__} (first use, " + "or the manager changed since its last bind) and CUDA graph capture is active; " + "an eager step must take and check its views first" + ) + _require( + cm, + ("get_index_k_buffer", "get_qsa_position_buffer", "get_buffers"), + "reads the draft layer's QSA index plane, position plane and paged K/V", + ) + g = self._pool_geometry(md.tokens_per_block) + layer = self.mtp_layer_idx + index_cache = cm.get_index_k_buffer(layer) + position_cache = cm.get_qsa_position_buffer() + check_position_plane(g, position_cache) + kv = self._kv_planes(cm, g, layer, index_cache) + share_slot = None + if self.mtp_index_share: + # The metadata sizes its shared selection buffers from the same resolved knob, + # one row per sparse layer of whichever manager is bound. + assert md.sparse_metadata_params.mtp_index_share, ( + "this target resolved IndexShare on but the engine's QSA metadata did not" + ) + slots = cm.qsa_shared_topk_slots + check_shared_selection(md.qsa_shared_topk_indices, slots, layer, self.block_topk) + share_slot = slots[layer] + self._draft_bound = _DraftBound(cm, index_cache, position_cache, kv, share_slot) # -------------------------------------------------------------- forward --- @@ -1342,16 +1789,26 @@ def _select_target(self, attn_metadata) -> "PrefillTarget | DecodeTarget": draft step routes through here per call.""" return self._decode if attn_metadata.num_contexts == 0 else self._prefill - def forward( + def begin_step( self, attn_metadata: AttentionMetadata, - input_ids: torch.IntTensor | None = None, - position_ids: torch.IntTensor | None = None, + input_ids: torch.IntTensor | None, + position_ids: torch.IntTensor | None, inputs_embeds: torch.FloatTensor | None = None, lora_params: dict | None = None, - **kwargs, - ): - """Check what holds for every step, bind the cache manager, then route.""" + spec_metadata=None, + vision: tuple | None = None, + ) -> None: + """Check what holds for every step, bind the cache manager, choose the target and + let the prefill target prepare the step. + + The shell calls it before `forward`, on every step. It is separate because a + breakable prefill CUDA graph captures `forward` and replays it without calling it: + whatever a step reads from the host (metadata, host buffers, the manager) is read + here, where it runs eagerly on a replayed step too. `vision` is `(rows, + mm_token_indices)` when the step carries media: the tower's `[R, hidden]` rows and + the `[R]` packed-batch rows they replace; the prefill target splices them inside + its forward.""" assert self._prefill is not None and self._decode is not None, ( "load_weights must run before forward" ) @@ -1359,33 +1816,74 @@ def forward( assert isinstance(attn_metadata, TrtllmAttentionMetadata) # Unimplemented inputs must fail loudly, not be dropped. assert lora_params is None, "LoRA is not implemented by this target" - # `input_ids` is required even with `inputs_embeds` (upstream passes - # `input_ids=None` after splicing): the PLE block hashes token ids every step, - # and text gates would not catch a mis-hash. + # `input_ids` is required on every step: the PLE block hashes token ids whether or + # not vision rows replace some embeddings, and text gates would not catch a mis-hash. assert input_ids is not None, ( "this target's PLE layer hashes token ids, so input_ids is required " - "on every step, including one whose embeddings were spliced upstream" + "on every step, including one that carries vision rows" + ) + # The token embedding is this target's own and vision rows are spliced into it + # inside the forward; an embedding handed in from outside would be dropped on every + # replay of a breakable prefill graph. + assert inputs_embeds is None, ( + "this target computes its token embedding itself and splices vision rows into " + "it; an engine-supplied inputs_embeds is not implemented" ) # `spec_metadata` is present exactly when drafting. On MTP_EAGLE_ONE_MODEL it # sets `layers_to_capture = ()`, so the trunk owes no capture hook. - spec_metadata = kwargs.get("spec_metadata") assert (spec_metadata is not None) == self.mtp_enabled, ( f"spec_metadata {'arrived' if spec_metadata is not None else 'is absent'} " f"but the model was built with mtp_enabled={self.mtp_enabled}; the " "draft path and the weights it loads are decided at construction" ) - # Keyed on the manager, not a first-forward flag: capacity estimation swaps - # the manager under the same module (see `_bind_cache_manager`). + # Keyed on the manager, not a first-forward flag: capacity estimation swaps the + # manager under the same module (see `_bind_cache_manager`). if self._bound is None or attn_metadata.kv_cache_manager is not self._bound.cm: self._bind_cache_manager(attn_metadata, position_ids) + target = self._select_target(attn_metadata) + if target is self._prefill: + target.prepare(attn_metadata, input_ids, position_ids, vision, spec_metadata) + else: + assert vision is None, ( + "a generation-only step carries vision rows; the executor attaches media " + "to context requests only" + ) + self._step_begun = True + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: torch.IntTensor | None = None, + position_ids: torch.IntTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + lora_params: dict | None = None, + **kwargs, + ): + """Run the step `begin_step` prepared on the target it chose. + + Under a breakable prefill CUDA graph this is the body the engine captures for a + step with context rows; a replayed step never enters it, which is why nothing here + may read the step from the host. Returns the logits input, or under speculation the + decode target's `(final, hc)` pair and the prefill target's one packed tensor + (`split_spec_output`).""" + assert self._step_begun, ( + "the shell calls begin_step before every forward; this forward has no prepared step" + ) + self._step_begun = False return self._select_target(attn_metadata).forward( - attn_metadata, - input_ids=input_ids, - position_ids=position_ids, - inputs_embeds=inputs_embeds, - **kwargs, + attn_metadata, input_ids=input_ids, position_ids=position_ids, **kwargs ) + def split_spec_output(self, output) -> tuple[torch.Tensor, torch.Tensor]: + """`(final, hc)` out of what `forward` returned under speculation: `final` is the + `[T, hidden]` mix `lm_head` reads and `hc` the `[T, hc_count * hidden]` stream the + MTP module reads. The decode target returns the pair; the prefill target returns one + `[T, hidden + stream]` tensor, of which the two are column views.""" + if isinstance(output, tuple): + return output + final, hc = torch.split(output, [self.hidden, self.stream], -1) + return final, hc + class _LayerWeights(NamedTuple): """One trunk layer's weights as the views the ops read (GEMM operands are `.t()` @@ -1423,9 +1921,14 @@ class _LayerWeights(NamedTuple): fc1_s: torch.Tensor fc2_w: torch.Tensor fc2_s: torch.Tensor - out1_scale: torch.Tensor - out1_gate_scale: torch.Tensor - out2_scale: torch.Tensor + # trtllm-gen scalars; None on CuteDSL. + out1_scale: torch.Tensor | None + out1_gate_scale: torch.Tensor | None + out2_scale: torch.Tensor | None + # CuteDSL scalars; None on trtllm-gen. + alpha1: torch.Tensor | None + alpha2: torch.Tensor | None + global_sf: torch.Tensor | None sh_gu: torch.Tensor sh_dn: torch.Tensor sh_gate: torch.Tensor @@ -1434,6 +1937,12 @@ class _LayerWeights(NamedTuple): def _layer_views(core: Qwen4ExpTp1ModelingV2) -> list[_LayerWeights]: """Every trunk layer's `_LayerWeights`, read off the realized weights.""" w = core.w + # The expert stacks are uint8 storage; each backend reads the packed nibbles and the + # block scales as its own dtypes. + if core.trunk_moe_backend == "CUTEDSL": + w_dtype, s_dtype = torch.float4_e2m1fn_x2, torch.uint8 + else: + w_dtype, s_dtype = torch.uint8, torch.float8_e4m3fn layers = [] for i in range(core.num_layers): attn: dict = dict( @@ -1499,13 +2008,16 @@ def _layer_views(core: Qwen4ExpTp1ModelingV2) -> list[_LayerWeights]: **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"], + fc1_w=w[f"l{i}_fc1_w"].view(w_dtype), + fc1_s=w[f"l{i}_fc1_s"].view(s_dtype), + fc2_w=w[f"l{i}_fc2_w"].view(w_dtype), + fc2_s=w[f"l{i}_fc2_s"].view(s_dtype), + out1_scale=scalars.get("output1_scale_scalar"), + out1_gate_scale=scalars.get("output1_scale_gate_scalar"), + out2_scale=scalars.get("output2_scale_scalar"), + alpha1=scalars.get("alpha1"), + alpha2=scalars.get("alpha2"), + global_sf=scalars.get("global_sf"), 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(), @@ -1574,6 +2086,9 @@ class _MtpWeights(NamedTuple): o: torch.Tensor q_norm: torch.Tensor k_norm: torch.Tensor + idx: torch.Tensor + idx_qn: torch.Tensor + idx_kn: torch.Tensor router: torch.Tensor fc1_w: torch.Tensor fc1_s: torch.Tensor @@ -1609,6 +2124,9 @@ def _mtp_views(core: Qwen4ExpTp1ModelingV2) -> _MtpWeights: o=w[f"{p}o_w"].t(), q_norm=w[f"{p}q_norm"], k_norm=w[f"{p}k_norm"], + idx=w[f"{p}idx_w"].t(), + idx_qn=w[f"{p}idx_qn"], + idx_kn=w[f"{p}idx_kn"], router=w[f"{p}router"].t(), fc1_w=w[f"{p}fc1_w"], fc1_s=w[f"{p}fc1_s"], @@ -1623,6 +2141,230 @@ def _mtp_views(core: Qwen4ExpTp1ModelingV2) -> _MtpWeights: ) +def _select(core, scores, row_starts, row_ends, topk_indices, live_cols=None) -> None: + """Select each row's top-`block_topk` groups into `topk_indices` with `qsa_topk_select`, + then rewrite the row canonically, in place. + + The radix engine compacts its picks through a shared-memory atomic, so the order is + launch-dependent, and where the tie class at the budget's last score straddles the + budget -- every zero-score group of a row with fewer positive scores than the budget, + on this checkpoint's relu scores -- which tied groups are kept can be too. The + expansion keeps the row's order and the sparse GQA reduces in it, so a raw row past the + dense threshold is not reproducible run to run. `qsa_canonical_selection` rewrites it + into the window's canonical selection -- the groups strictly above the budget's last + score, then the highest-index (most recent) ties, ascending, `-1` last -- a function of + the scores alone, and the engine's own set in ascending order wherever the budget's + last score is untied. A canonical row is prefix-packed, which the expansion and the + IndexShare capture and reuse require. + + `index_topk` is the block budget: the entry's default (2048) is 4x the buffer width and + the op ignores the row stride, so it mis-writes silently. + + The rewrite reads every column it is given, and `scores` is as wide as the block table, + not the live window. A caller that knows a host-side bound on the rows' windows passes + it as `live_cols`: the rewrite then reads only those columns and returns the same rows.""" + qsa_topk_select(scores, row_starts, row_ends, topk_indices, index_topk=core.block_topk) + if live_cols is not None: + scores = scores[:, :live_cols] + qsa_canonical_selection(scores, row_starts, row_ends, topk_indices) + + +def _store_kv(core, kv, k, v, request_indices, logical_positions, block_table, tokens_per_block): + """Append this step's K/V rows to the layer's paged pool: a plain move over a bf16 or fp8 + pool (an fp8 pool takes e4m3 unscaled, as the dense op does at the unit KV scales), the + quantizing store over an NVFP4 one, at the same second-level scales as the dense op.""" + if kv.k_scale is None: + qsa_paged_kv_store( + k, v, kv.k, kv.v, request_indices, logical_positions, block_table, tokens_per_block + ) + return + qsa_paged_kv_store_nvfp4( + k, + v, + kv.k, + kv.v, + kv.k_scale, + kv.v_scale, + request_indices, + logical_positions, + block_table, + tokens_per_block, + core.kv_scale_orig_quant, + ) + + +def _sparse_gqa( + core, + kv, + q, + block_table, + selected_tokens, + request_indices, + tokens_per_block, + query_positions, + compress_ratio, +): + """Sparse GQA over the layer's paged pool (dequantizing an NVFP4 one). The prefill pair + `query_positions` / `compress_ratio` is a bound on a context row's scan, or None, None.""" + if kv.k_scale is None: + return qsa_paged_sparse_gqa( + q=q, + k_cache=kv.k, + v_cache=kv.v, + block_table=block_table, + selected_tokens=selected_tokens, + request_indices=request_indices, + tokens_per_block=tokens_per_block, + softmax_scale=core.softmax_scale, + query_positions=query_positions, + compress_ratio=compress_ratio, + ) + return qsa_paged_sparse_gqa_nvfp4( + q=q, + k_cache=kv.k, + v_cache=kv.v, + k_scale_cache=kv.k_scale, + v_scale_cache=kv.v_scale, + kv_scale_quant_orig=core.kv_scale_quant_orig, + block_table=block_table, + selected_tokens=selected_tokens, + request_indices=request_indices, + tokens_per_block=tokens_per_block, + softmax_scale=core.softmax_scale, + query_positions=query_positions, + compress_ratio=compress_ratio, + ) + + +def _live_score_columns(core, md, batch: int) -> int: + """A host-side bound on every row's selection window in a step with context rows: the + longest sequence's compressed length, with the verify width as slack for the draft + rows that sit past the host's KV lengths.""" + longest = int(md.kv_lens_runtime[:batch].max()) if batch > 0 else 0 + return (longest + _MAX_VERIFY_WIDTH) // core.compress_ratio + 1 + + +def _replay_verify(core, mq_out, a_d, b_d, ssm_states, i, L, si_d, n, p, replay) -> torch.Tensor: + """One GDN layer's verify recurrence under the cached replay: `n` generation sequences + of `p` rows over the conv output `mq_out` (`[n * p, conv_dim]`, token-major) and the + raw `a`/`b` projections, with the fp32 gating parameters fused into the kernel. + + `replay` is `(work_items, n_writes, partitioned)`; the work items are given exactly + when the manager runs its all-layer commit after the step (without that commit every + history wrap goes uncommitted, with it a wrap committed here is committed twice). + `mq_out` is the conv's fresh, contiguous output, so the kernel reads q/k/v from it + directly (`packed_qkv`) and never from the views. The conv before it does not signal + dependents, so the kernel launches without PDL.""" + work_items, n_writes, partitioned = replay + state = core._bound.replay + old_x, old_B, old_dt, old_dA = state.old[i] + qv, kv, vv = torch.split(mq_out, [core.key_dim, core.key_dim, core.value_dim], 1) + out = fused_recurrent_gated_delta_rule_cached_replay_update( + q=torch.reshape(qv, [n, p, core.num_k_heads, core.head_k_dim]), + k=torch.reshape(kv, [n, p, core.num_k_heads, core.head_k_dim]), + v=torch.reshape(vv, [n, p, core.num_v_heads, core.head_v_dim]), + g=torch.reshape(a_d, [n, p, core.num_v_heads]), + beta=torch.reshape(b_d, [n, p, core.num_v_heads]), + ssm_states=ssm_states, + state_indices=si_d, + old_u=old_x, + old_k=old_B, + old_G=old_dt, + old_beta=old_dA, + cache_buf_idx=state.cache_buf_idx, + prev_num_accepted_tokens=state.prev_num_accepted, + history_size=state.history, + scale=None, + use_qk_l2norm_in_kernel=True, + A_log=L.a_log, + dt_bias=L.dt_bias, + launch_with_pdl=False, + replay_work_items=work_items, + n_writes=n_writes, + block_v=None, + num_warps=None, + use_l2_state_cache=None, + use_l2_streaming_inputs=None, + use_l2_shared_inputs=None, + main_num_stages=None, + packed_qkv=mq_out, + use_all_layer_commit=partitioned, + output=None, + ) + return torch.reshape(out, [1, n * p, core.num_v_heads, core.head_v_dim]) + + +def _replay_work(core, mm, num_decodes: int) -> tuple: + """`_replay_verify`'s `replay` argument for a verify step of `num_decodes` generation + requests: the partitioned replay exactly when `update_mamba_states` runs the manager's + all-layer commit after the step.""" + partitioned = ( + core._bound.replay.all_layer_commit + and num_decodes >= CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE + ) + return ( + mm.replay_work_items[:num_decodes] if partitioned else None, + mm.replay_n_writes if partitioned else None, + partitioned, + ) + + +class _KvRangeProbe: + """`TRTLLM_MODELING_V2_VALIDATE` over an fp8 KV pool: reports, once, the largest K/V + magnitude the first eager forward's attention layers store. + + The pool takes K/V at the unit scale, so anything past the e4m3 maximum saturates + silently (a bf16 inf included). One host sync per attention layer of that one forward; + a forward under CUDA graph capture is passed over without a read.""" + + def __init__(self, last_layer: int) -> None: + self._last = last_layer + self._worst = 0.0 + + def observe(self, layer: int, k: torch.Tensor, v: torch.Tensor) -> bool: + """Fold one attention layer's K/V in; False once the forward's last layer has + reported.""" + if torch.cuda.is_current_stream_capturing(): + return True + peak = float(torch.maximum(k.abs().amax(), v.abs().amax())) + # A NaN stays: `max` would drop it. + self._worst = math.nan if math.isnan(peak + self._worst) else max(self._worst, peak) + if layer != self._last: + return True + limit = torch.finfo(torch.float8_e4m3fn).max + # Written so that NaN fails too. + if not self._worst <= limit: + logger.warning( + f"FP8 KV cache: the first forward stores a K/V magnitude of {self._worst:g}, " + f"past the e4m3 maximum {limit:g}; the pool is written at the unit scale, " + "so those values saturate" + ) + return False + + +class _PrefillStep(NamedTuple): + """One step's host-side state, built by `PrefillTarget.prepare` and read by the forward's + breaks. + + `num_tokens` is the batch's row count `n` and `physical_tokens` the `T` the forward runs + over (`n`, or the prefill CUDA graph bucket it was padded to); `input_ids` is the batch's + `n` ids. `attn` is thop_attention's per-step arguments; `gdn`, `ple` and `qsa` are the + recurrent, PLE and QSA bookkeeping tuples their breaks unpack, each cut to the batch's + rows; `stage` says the speculative rollback stages this verify step; `vision` is the + shell's `(rows, mm_token_indices)` or None. A new one replaces it every step, so no + break reads a previous step's.""" + + num_tokens: int + physical_tokens: int + input_ids: torch.Tensor + attn: dict + gdn: tuple + ple: tuple + qsa: tuple + stage: bool + vision: tuple | None + + class PrefillTarget: """Any batch with context rows (`num_contexts > 0`), possibly with generation rows beside them; the recurrent layers' decode-only branches never run here.""" @@ -1633,8 +2375,16 @@ def __init__(self, core: Qwen4ExpTp1ModelingV2) -> None: # Off unless TRTLLM_MODELING_V2_VALIDATE asks; each target owns its flag (the # two `step_args` projections differ). self._contract_pending = step_contract_enabled() + # Dropped once it has reported. + self._kv_probe = ( + _KvRangeProbe(core.attn_layers[-1]) + if self._contract_pending and core.quant_mode == QM_FP8_KV_CACHE + else None + ) self._layers = _layer_views(core) self._stack = _stack_views(core) + # The step `prepare` built, read by the breaks. + self._step: _PrefillStep | None = None def _check_step_contract(self, md) -> None: """Opt-in first-forward fail-fast (TRTLLM_MODELING_V2_VALIDATE): calling @@ -1706,75 +2456,132 @@ def _grouped_norm(self, x: torch.Tensor, weights: tuple) -> torch.Tensor: [flashinfer_gemma_rmsnorm(part, w, core.eps) for w, part in zip(weights, parts)], 1 ) - def _moe_block(self, x: torch.Tensor, L: _LayerWeights, dt: torch.dtype) -> torch.Tensor: + def _moe_block( + self, x: torch.Tensor, L: _LayerWeights, dt: torch.dtype, routed_out: torch.Tensor | None + ) -> 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.""" + plus a bf16 shared expert folded in by a per-token sigmoid gate. `routed_out` + is the CuteDSL path's `[T, hidden]` output buffer (None on trtllm-gen, whose + runner allocates its own).""" 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] + if core.trunk_moe_backend == "CUTEDSL": + # fp32 weights: the dtype the unpermute scales by. + topk_ids, topk_w = renorm_moe_routing_op( + logits, topk=core.topk, output_dtype=torch.float32 + ) + xq, xsf = fp4_quantize(x, L.moe_gscale, _SF_VEC, False, _SF_LINEAR) + # The sort pads each expert's (token, expert) pairs to whole tiles and emits + # the tile tables both GEMMs walk. + tile_expert, tile_limit, expanded_to_permuted, permuted_to_expanded, _, live_tiles = ( + moe_sort( + token_selected_experts=topk_ids, + token_final_scales=topk_w, + num_experts=core.num_experts, + top_k=core.topk, + local_expert_offset=0, + local_num_experts=core.num_experts, + tile_tokens_dim=_CUTEDSL_TILE_SIZE, + ) + ) + # FC1 gathers its rows from the unpermuted activation, applies SwiGLU and + # requantizes with the layer's one FC2 input scale. `fp4_quantize` returns the + # linear scales flat; the GEMM reads them as [T, hidden / 16]. + h, h_sf = cute_dsl_nvfp4_gather_grouped_gemm_act_fusion( + input=xq.view(torch.float4_e2m1fn_x2), + weight=L.fc1_w, + input_scale=xsf.view(x.shape[0], core.hidden // _SF_VEC), + weight_scale=L.fc1_s, + alpha=L.alpha1, + tile_idx_to_group_idx=tile_expert, + tile_idx_to_mn_limit=tile_limit, + permuted_idx_to_expanded_idx=permuted_to_expanded, + num_non_exiting_tiles=live_tiles, + global_sf=L.global_sf, + num_experts=core.num_experts, + top_k=core.topk, + num_local_experts=core.num_experts, + local_expert_offset=0, + tile_size=_CUTEDSL_TILE_SIZE, + scaling_vector_size=_SF_VEC, + activation_type=int(ActivationType.Swiglu), + ) + # FC2 writes one bf16 row per permuted row; the unpermute writes every row of + # `routed_out` as the routing-weighted sum of its token's top-k rows. + y = cute_dsl_nvfp4_grouped_gemm( + input=h, + weight=L.fc2_w, + input_scale=h_sf, + weight_scale=L.fc2_s, + alpha=L.alpha2, + tile_idx_to_group_idx=tile_expert, + num_non_exiting_tiles=live_tiles, + num_experts=core.num_experts, + top_k=core.topk, + num_local_experts=core.num_experts, + local_expert_offset=0, + tile_size=_CUTEDSL_TILE_SIZE, + output_dtype=torch.bfloat16, + scaling_vector_size=_SF_VEC, + ) + routed = moe_unpermute_inplace( + permuted_input=y, + output=routed_out, + expanded_idx_to_permuted_idx=expanded_to_permuted, + topk_scales=topk_w, + ) + else: + 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): + def _attention_block(self, x, i, L, qkv, index_scratch, attn_out, pos, pos3): """One gated GQA layer plus its QSA index branch. The q projection is twice as wide as q (q and output gate interleave per head, as HF chunks the last dim); only the q half is packed into the QKV - buffer and the gate half is read later through its own strides. The index - plane advances on every step, including below the threshold, so a request - growing past it finds its prefix compressed; above it the main K/V append - becomes this layer's own call.""" + buffer and the gate half is read later through its own strides. The + projections, the packing and both ropes are row-wise and run over all `T` rows + here; the cache writes, the selection and the attention core read the step's + batch and run in the `_attention_core` break, which leaves the attention output + in `attn_out`.""" 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) @@ -1814,32 +2621,15 @@ def _attention_block(self, x, i, L, qkv, index_scratch, attn_out, qsa, step): mrope_section1=core.mrope_section1, mrope_section2=core.mrope_section2, ) + if self._kv_probe is not None and not self._kv_probe.observe(i, k_dst, v_dst): + self._kv_probe = None - # ---- 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). + # One replicated projection produces [index_heads q | 1 k]; its raw key is what + # the break stores. The query's norm and rotation run here in the fused call's + # packed layout, on a scratch copy: 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. 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( @@ -1861,9 +2651,68 @@ def _attention_block(self, x, i, L, qkv, index_scratch, attn_out, qsa, step): mrope_section1=(_QSA_MROPE_SECTION or (0, 0, 0))[1], mrope_section2=(_QSA_MROPE_SECTION or (0, 0, 0))[2], ) + eager_on_graph(qkv, qk_index, index_scratch, attn_out, i, inner=self._attention_core) + fused_sigmoid_mul_inplace(attn_out, gate) + return cublas_mm(attn_out, L.o) + + def _attention_core(self, qkv, qk_index, index_scratch, attn_out, i) -> None: + """Break: the index plane's store and compression, the dense/sparse decision and the + attention core over the step's `n` rows. Writes rows `[:n]` of `attn_out` + `[T, q_width]` in place and zeroes rows `[n:]`. + + The index plane advances on every step, including below the threshold, so a + request growing past it finds its prefix compressed; above it the main K/V append + becomes this layer's own call. The raw index key is stored before any + norm/rotation: the compression step applies both itself, so a pre-normalized key + is double-normalized (invisible below the sparse threshold). + + On a verify step (`stage`) the generation rows' index slots (and, on the first + attention layer, their position slots) are snapshotted for `SpeculativeRollback` + before the store and compression overwrite them.""" + core = self.core + L = self._layers[i] + step = self._step + ( + req_idx, + logical, + seq_lens_q, + visible_blocks, + topk_indices, + topk_row_starts, + block_table, + tokens_per_block, + coords, + sparse, + live_cols, + ) = step.qsa + num_tokens = step.num_tokens + qkv = qkv[:num_tokens] + 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) + + index_cache = core._bound.index_cache[i] + token_k = torch.split( + torch.reshape( + qk_index[:num_tokens], [num_tokens, core.index_heads + core.index_kv_heads, -1] + ), + [core.index_heads, core.index_kv_heads], + 1, + )[1] + if step.stage: + core.spec_rollback.stage_index(index_cache, core._bound.position_cache) + qsa_paged_index_store( + token_k, + coords, + index_cache, + core._bound.position_cache, + req_idx, + logical, + block_table, + tokens_per_block, + ) q_index = torch.split( torch.reshape( - scratch_h, + index_scratch[:num_tokens], [num_tokens, core.index_heads + 2 * core.index_kv_heads, core.index_head_dim], ), [core.index_heads, core.index_kv_heads, core.index_kv_heads], @@ -1884,12 +2733,13 @@ def _attention_block(self, x, i, L, qkv, index_scratch, attn_out, qsa, step): ) if sparse: - k_cache, v_cache = core._bound.kv[i] + kv = core._bound.kv[i] # Append the fused-norm+rope'd K/V here: nothing else writes the pool on - # this path, and selection reads the current token from it. - qsa_paged_kv_store( - k_dst, v_dst, k_cache, v_cache, req_idx, logical, block_table, tokens_per_block - ) + # this path, and selection reads the current token from it. Over an fp8 pool the + # store and the sparse GQA below take e4m3 unscaled, as the dense op does at the + # unit KV scales; over an NVFP4 pool they quantize and dequantize at the dense op's + # second-level scales. + _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block) scores = qsa_paged_index_scores( q=q_index, index_cache=index_cache, @@ -1901,11 +2751,7 @@ def _attention_block(self, x, i, L, qkv, index_scratch, attn_out, qsa, step): 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 - ) + _select(core, scores, topk_row_starts, visible_blocks, topk_indices, live_cols) selected = qsa_expand_block_indices( topk_indices, logical, @@ -1913,50 +2759,67 @@ def _attention_block(self, x, i, L, qkv, index_scratch, attn_out, qsa, step): 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], + attn_out[:num_tokens].copy_( + torch.reshape( + _sparse_gqa( + core, + kv, + q=q_dst, + block_table=block_table, + selected_tokens=selected, + request_indices=req_idx, + tokens_per_block=tokens_per_block, + # 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( + thop_attention( q=qkv, - output=attn_out, + output=attn_out[:num_tokens], local_layer_idx=core._bound.pool_idx[i], num_heads=core.heads_q, num_kv_heads=core.heads_kv, head_size=core.head_dim, quant_mode=core.quant_mode, - **step, + kv_scale_orig_quant=core.kv_scale_orig_quant, + kv_scale_quant_orig=core.kv_scale_quant_orig, + **step.attn, **_CALL_CONSTANTS, ) - fused_sigmoid_mul_inplace(out, gate) - return cublas_mm(out, L.o) + attn_out[num_tokens:].zero_() - def _gdn_verify(self, mq_d, a_d, b_d, ssm_states, i, L, si_d, int_idx, n, p) -> torch.Tensor: + def _gdn_verify( + self, mq_d, a_d, b_d, ssm_states, i, L, si_d, int_idx, n, p, replay + ) -> torch.Tensor: """The recurrent half of a speculative verify step: `n` generation sequences of - `p = runtime_draft_len + 1` tokens, with per-step conv window and recurrent - state written to the runtime's rollback buffers, not the pools. - - The pools must not move: which prefix is real is unknown until the sampler - accepts, and `update_mamba_states` then promotes step `k - 1`. Advancing them - would leave rejected tokens' state in every recurrent layer, invisible to the - acceptance rate. `causal_conv1d_triton_update` is the only conv with per-step - snapshots (its result for a pad-sentinel slot is uninitialized, safe because - every row is a real sequence); `pack_gdn_decode_qkv` does no Q/K L2 norm, - hence `use_qk_l2norm_in_kernel=True` below.""" + `p = runtime_draft_len + 1` tokens, written so that the accepted prefix, unknown + until the sampler has run, can be kept and the rest dropped. + + The conv writes each row's window to the runtime's per-step buffer and leaves the + conv pool alone; `update_mamba_states` promotes window `k - 1`. + `causal_conv1d_triton_update` is the only conv with per-step snapshots (its result + for a pad-sentinel slot is uninitialized, safe because every row is a real + sequence). The recurrence has two forms, bound per manager + (`_bind_step_buffers`): + + * cached replay (`replay` given): the pool row is a checkpoint; one call rebuilds + the current state from it and the slot's pending history, runs the `p` rows and + appends them to the history. Rejected rows cost nothing to undo, because the + manager advances the pending count by the accepted tokens only. A full history + advances the checkpoint inside the kernel below + `CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE` requests and in the manager's + all-layer commit from there on, hence `replay` carries the work items exactly + under the manager's own condition. + * legacy (`replay` None): every row's recurrent state goes to the per-step buffer + instead of the pool, and the promotion copies step `k - 1` into it. Advancing + the pool directly would leave rejected tokens' state in every recurrent layer, + invisible to the acceptance rate. `pack_gdn_decode_qkv` does no Q/K L2 norm, + hence `use_qk_l2norm_in_kernel=True` below.""" core = self.core # [n*p, C] -> [n, C, p]: the multi-token conv reads channel-major (the # transposed view is certified bitwise equal to the contiguous form). @@ -1977,6 +2840,8 @@ def _gdn_verify(self, mq_d, a_d, b_d, ssm_states, i, L, si_d, int_idx, n, p) -> ) # Back to token-major (a view). mq_out = torch.reshape(torch.transpose(conv_out, 1, 2), [n * p, core.conv_dim]) + if replay is not None: + return _replay_verify(core, mq_out, a_d, b_d, ssm_states, i, L, si_d, n, p, replay) q, k, v = pack_gdn_decode_qkv( mq_out, core.num_k_heads, core.head_k_dim, core.num_v_heads, core.head_v_dim ) @@ -2017,10 +2882,27 @@ def _gdn_verify(self, mq_d, a_d, b_d, ssm_states, i, L, si_d, int_idx, n, p) -> ) return torch.reshape(out, [1, n * p, core.num_v_heads, core.head_v_dim]) - def _gdn_block(self, x, i, L, gdn_state) -> torch.Tensor: - """One Gated-DeltaNet layer over a batch with a context prefix. Every branch - is on a host int the engine fixes per capture, never a tensor value.""" + def _gdn_block(self, x, i, L, core_out) -> torch.Tensor: + """One Gated-DeltaNet layer over a batch with a context prefix: the four input + projections and the gated-norm epilogue run over all `T` rows here, the + recurrence in the `_gdn_core` break, which leaves its output in `core_out` + `[1, T, num_v_heads, head_v_dim]`.""" + 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) + eager_on_graph(mixed_qkv, a, b, core_out, i, inner=self._gdn_core) + return self._gdn_epilogue(core_out, z, L, x.shape[0]) + + def _gdn_core(self, mixed_qkv, a, b, out, i) -> None: + """Break: the recurrence over the step's `n` rows (the conv, the post-conv gating and + the chunked delta rule, or under speculation the verify path for the generation + rows). Writes rows `[:n]` of `out` `[1, T, num_v_heads, head_v_dim]` in place and + zeroes rows `[n:]`. Every branch is on a host int of the step, never a tensor + value.""" core = self.core + L = self._layers[i] + step = self._step ( num_prefill_tokens, num_decode_tokens, @@ -2037,13 +2919,12 @@ def _gdn_block(self, x, i, L, gdn_state) -> torch.Tensor: gen_p, int_idx, qsl_long_p, - ) = gdn_state - num_tokens = x.shape[0] - - 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) + replay, + ) = step.gdn + num_tokens = step.num_tokens + mixed_qkv = mixed_qkv[:num_tokens] + a = a[:num_tokens] + b = b[:num_tokens] conv_states = core._bound.conv[i] ssm_states = core._bound.ssm[i] @@ -2104,9 +2985,12 @@ def _gdn_block(self, x, i, L, gdn_state) -> torch.Tensor: 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 + mq_d, a_d, b_d, ssm_states, i, L, si_d, int_idx, num_decodes, gen_p, replay ) - return self._gdn_epilogue(torch.cat([out_p, out_d], 1), z, L, num_tokens) + out[:, :num_prefill_tokens].copy_(out_p) + out[:, num_prefill_tokens:num_tokens].copy_(out_d) + out[:, num_tokens:].zero_() + return causal_conv1d_update( x=mq_d, conv_state=conv_states, @@ -2160,7 +3044,7 @@ def _gdn_block(self, x, i, L, gdn_state) -> torch.Tensor: softplus_threshold=_SOFTPLUS_THRESHOLD, ) # q and k are already L2-normalized; the recurrence must not repeat it. - out, _ = chunk_gated_delta_rule( + recurrence, _ = chunk_gated_delta_rule( q, k, v, @@ -2175,7 +3059,8 @@ def _gdn_block(self, x, i, L, gdn_state) -> torch.Tensor: use_qk_l2norm_in_kernel=False, output=None, ) - return self._gdn_epilogue(out, z, L, num_tokens) + out[:, :num_tokens].copy_(recurrence) + out[:, num_tokens:].zero_() def _gdn_epilogue(self, out, z, L, num_tokens) -> torch.Tensor: """Gated RMS norm and output projection shared by every recurrence.""" @@ -2191,9 +3076,10 @@ def _gdn_epilogue(self, out, z, L, num_tokens) -> torch.Tensor: ) return cublas_mm(torch.reshape(o, [num_tokens, core.value_dim]), L.out) - def _ple_block(self, hc, input_ids, ple_state) -> torch.Tensor: - """The n-gram side path on one decoder layer; the caller adds the result into - the hyper-connection stream. + def _ple_block(self, hc, out) -> None: + """Break: the n-gram side path on one decoder layer, over the step's `n` rows. Writes + rows `[:n]` of `out` `[T, stream]` in place and zeroes rows `[n:]`; the caller adds + `out` 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 @@ -2201,11 +3087,23 @@ def _ple_block(self, hc, input_ids, ple_state) -> torch.Tensor: scale; key/value projections and the gate scale one value into each stream; a dilated depthwise short conv over a paged state is added back. A sequence carrying nothing is seeded with `eos_token_id`, not 0: id 0 is a real token - the hash cannot distinguish.""" + the hash cannot distinguish. + + The id history and the conv state publish the state after the whole chunk. On + a verify step (`stage`) the generation rows' pre-step pool rows, id windows and + conv input go to `SpeculativeRollback`, which rewrites both pools to the + accepted prefix once the sampler has decided it.""" core = self.core S = self._stack - conv_states, ngram_pool, si_all, carry_all, qsl = ple_state - + step = self._step + conv_states, ngram_pool, si_all, carry_all, qsl = step.ple + stage = step.stage + num_tokens = step.num_tokens + input_ids = step.input_ids + hc = hc[:num_tokens] + + if stage: + core.spec_rollback.stage_ple_pools(ngram_pool, conv_states) window = ngram_token_history( input_ids, ngram_pool, si_all, carry_all, qsl, core.eos_token_id, core.ngram_size ) @@ -2231,9 +3129,12 @@ def _ple_block(self, hc, input_ids, ple_state) -> torch.Tensor: gated = gated_ngram_injection( key_normed, query_normed, value, core.hc_count, _PLE_CLAMP_MIN ) + conv_in = self._grouped_norm(gated, S.ple_nc) + if stage: + core.spec_rollback.stage_ple_rows(window, conv_in) # The short conv reads channel-major; a bare transpose is the certified form. convolved = dilated_short_conv_paged( - x=torch.transpose(self._grouped_norm(gated, S.ple_nc), 0, 1), + x=torch.transpose(conv_in, 0, 1), weight=S.ple_conv_w, bias=None, conv_states=conv_states, @@ -2244,72 +3145,67 @@ def _ple_block(self, hc, input_ids, ple_state) -> torch.Tensor: silu_activation=True, pad_slot_id=_PAD_SLOT_ID, ) - return torch.add(gated, torch.transpose(convolved, 0, 1)) + torch.add(gated, torch.transpose(convolved, 0, 1), out=out[:num_tokens]) + out[num_tokens:].zero_() + + def _splice(self, h) -> None: + """Break: write the step's vision rows over their placeholder rows of the token + embedding `h` `[T, hidden]` in place, and zero rows `[n:]`. + + `mm_token_indices` names every placeholder row of the packed batch in the order the + per-request rows concatenate in. An indexed write, not a slice copy: two images in + one prompt are separated by framing tokens, so the placeholder rows are several + disjoint runs. A text step writes nothing but the padding. The rows are the step's: + a replayed graph never sees the previous step's.""" + step = self._step + if step.vision is not None: + rows, mm_token_indices = step.vision + h.index_copy_(0, mm_token_indices, rows) + h[step.num_tokens :].zero_() # -------------------------------------------------------------- forward --- - def forward( + def prepare( self, attn_metadata: AttentionMetadata, - input_ids: torch.IntTensor | None = None, - position_ids: torch.IntTensor | None = None, - inputs_embeds: torch.FloatTensor | None = None, - **kwargs, - ): + input_ids: torch.IntTensor, + position_ids: torch.IntTensor, + vision: tuple | None, + spec_metadata, + ) -> None: + """Everything this step reads from the host, kept for the forward's breaks as one + `_PrefillStep`: the contract check, the attention op's per-step arguments and the + batch's bookkeeping, built once and shared by every layer. + + Called by the core's `begin_step` on every step with context rows, before the + forward. On a step a breakable prefill graph replays, this and the breaks are all of + this target that runs.""" core = self.core md = attn_metadata self._check_step_contract(md) - step = self.step_args(md) - spec_metadata = kwargs.get("spec_metadata") + attn_args = self.step_args(md) mm = md.mamba_metadata - # Three views of one `[3, 1, T]` buffer: `pos` (row 0, flat `[T]`: QSA index - # rope), `pos3` (contiguous `[3, T]` for `use_mrope=True`; the engine hands a - # slice of a larger buffer and the op rejects the row stride, hence the - # copy), `coords` (`[T, 3]`, stored by the QSA position side cache). - pos = torch.reshape(torch.split(position_ids, [1, _QSA_POSITION_AXES - 1], 0)[0], [-1]) - coords = torch.transpose(torch.reshape(position_ids, [_QSA_POSITION_AXES, -1]), 0, 1) - pos3 = torch.empty( - [_QSA_POSITION_AXES, position_ids.shape[-1]], - dtype=torch.int32, - device=position_ids.device, + # **The activation rows are the batch's rows, or a prefill CUDA graph bucket's.** A + # breakable prefill graph pads the step's rows up to `padded_num_tokens` (zero ids, + # zero positions), and so does a padded step that runs eagerly; the forward runs over + # all of them and every break slices the batch's own rows back out. Anything else is + # an activation buffer longer than the batch, which would push padding rows through + # the recurrences into the state pools. + physical_tokens = input_ids.shape[0] + assert physical_tokens in (md.num_tokens, md.padded_num_tokens), ( + f"the engine handed {physical_tokens} activation rows for a batch its " + f"metadata describes as {md.num_tokens} tokens " + f"(padded_num_tokens={md.padded_num_tokens})" ) - pos3.copy_(torch.reshape(position_ids, [_QSA_POSITION_AXES, -1])) - - # Spliced vision rows are already the embedding; a lookup would overwrite them. - if inputs_embeds is None: - h = nn.functional.embedding(input_ids, core.w["embed"]) - else: - assert inputs_embeds.shape[0] == input_ids.shape[0], ( - f"the spliced embedding carries {inputs_embeds.shape[0]} rows for " - f"{input_ids.shape[0]} token ids" - ) - assert inputs_embeds.shape[1] == core.hidden, ( - f"the spliced embedding is {inputs_embeds.shape[1]} wide, not {core.hidden}" - ) - h = inputs_embeds - num_tokens = h.shape[0] - dt = h.dtype - dev = h.device # The batch's bookkeeping, built once and shared by every layer. + num_tokens = md.num_tokens batch = md.num_seqs num_prefills = md.num_contexts num_decodes = batch - num_prefills num_prefill_tokens = md.num_ctx_tokens num_decode_tokens = num_tokens - num_prefill_tokens - # Host-int check (graph-safe): an activation buffer longer than the batch - # would push padding rows through the recurrences into the state pools. - assert num_tokens == md.num_tokens, ( - f"the engine handed {num_tokens} activation rows for a batch its " - f"metadata describes as {md.num_tokens} tokens " - f"(padded_num_tokens={getattr(md, 'padded_num_tokens', None)}). A " - "prefill/piecewise CUDA-graph backend pads the activation buffer " - "past the batch and expects the model to slice it back; this " - "target does not implement that configuration. Decode CUDA graphs " - "pad with real dummy requests instead, so they are in the metadata " - "and this assert is silent for them." - ) # Rows per generation sequence: 1 without speculation, else # `runtime_draft_len + 1`; read from the metadata because the worker rewrites # it between draft steps. @@ -2342,12 +3238,10 @@ def forward( "sized from the second and written from the first" ) # Per-iteration counterpart of the core's `max_draft_len` bound. - assert 1 <= gen_p <= 4, ( - f"predicted_tokens_per_seq {gen_p} is outside the 1..4 range " - f"attention/thop_attention certifies" - ) + check_verify_width(gen_p) + check_replay_width(core._bound.replay, gen_p) is_verify = spec_metadata is not None and num_decodes > 0 and gen_p > 1 - step["predicted_tokens_per_seq"] = gen_p + attn_args["predicted_tokens_per_seq"] = gen_p si_all = mm.state_indices[:batch] hi_all = mm.has_initial_states[:batch] @@ -2359,11 +3253,14 @@ def forward( qsl_long = mm.query_start_loc_long int_idx = None qsl_long_p = None + replay = None if is_verify: # The runtime's arange over the slot pool, from the cache manager (what # `update_mamba_states` reads). assert core._bound.int_idx is not None, "the cache-manager bind has not run" int_idx = core._bound.int_idx[:num_decodes] + if core._bound.replay is not None: + replay = _replay_work(core, mm, num_decodes) # Context rows keep an ordinary chunked recurrence over their prefix. qsl_long_p = qsl_long[: num_prefills + 1] gdn_state = ( @@ -2382,6 +3279,7 @@ def forward( gen_p, int_idx, qsl_long_p, + replay, ) # PLE carry flags must cover the whole batch, but `has_initial_states` is @@ -2403,6 +3301,9 @@ def forward( # short/long decode graph families. sparse = batch > 0 and int(md.kv_lens_runtime[:batch].max()) > core.dense_threshold + # The `[T, 3]` transpose of the `[3, 1, T]` position buffer, which the QSA position + # plane stores. + coords = torch.transpose(torch.reshape(position_ids, [_QSA_POSITION_AXES, -1]), 0, 1) qsa = ( md.qsa_req_idx_per_token[:num_tokens], md.qsa_logical_positions[:num_tokens], @@ -2412,47 +3313,140 @@ def forward( md.qsa_topk_row_starts[:num_tokens], md.qsa_block_table, md.tokens_per_block, - pos, - pos3, - coords, + coords[:num_tokens], sparse, + _live_score_columns(core, md, batch), + ) + + # A verify step's generation rows write the PLE pools and the QSA side caches for + # every draft row, accepted or not; `SpeculativeRollback` keeps what its commit + # needs to undo the rejected ones. Reopened on every forward, so a step whose + # commit never ran (a failed forward) leaves nothing for the next one. Context + # rows are never staged. + stage = is_verify + rollback = core.spec_rollback + if rollback is not None: + rollback.discard() + if stage: + rollback.begin( + num_contexts=num_prefills, + prefix_tokens=num_prefill_tokens, + num_gens=num_decodes, + rows_per_gen=gen_p, + slots=si_d, + request_indices=md.qsa_req_idx_per_token[num_prefill_tokens:num_tokens], + logical_positions=md.qsa_logical_positions[num_prefill_tokens:num_tokens], + block_table=md.qsa_block_table, + tokens_per_block=md.tokens_per_block, + ) + + self._step = _PrefillStep( + num_tokens=num_tokens, + physical_tokens=physical_tokens, + input_ids=input_ids[:num_tokens], + attn=attn_args, + gdn=gdn_state, + ple=ple_state, + qsa=qsa, + stage=stage, + vision=vision, ) - # Scratch shared by all attention layers (each fully rewrites them). + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: torch.IntTensor, + position_ids: torch.IntTensor, + **kwargs, + ): + """The step `prepare` built, over the `T` activation rows the engine handed in: the + batch's `n` rows, or the prefill CUDA graph bucket they were padded to. + + **This is the body a breakable prefill CUDA graph captures**, once per bucket, and + replays for every step it pads to that bucket without running this Python. So + nothing here reads the step: every op is row-wise and sized by `T` alone, its inputs + are the engine's persistent `input_ids` / `position_ids` buffers, and every buffer is + allocated here. What depends on the step's batch runs in one of four breaks that + run eagerly between the captured segments and read `self._step`: `_splice` (once), + `_ple_block` (once), `_attention_core` (every attention layer) and `_gdn_core` (every + Gated-DeltaNet layer). Each writes rows `[:n]` of a fixed output allocated here and + zeroes rows `[n:]`, so the padding rows stay finite through the row-wise ops and + reach neither a state pool nor a real row. With the backend off every step is + eager and unpadded, and the breaks are plain calls.""" + core = self.core + num_tokens = input_ids.shape[0] + assert num_tokens == self._step.physical_tokens, ( + f"the forward got {num_tokens} activation rows and prepare saw " + f"{self._step.physical_tokens}" + ) + + # Two views of one `[3, 1, T]` buffer: `pos` (row 0, flat `[T]`: QSA index rope) and + # `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). + pos = torch.reshape(torch.split(position_ids, [1, _QSA_POSITION_AXES - 1], 0)[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])) + + # The token embedding with this step's vision rows spliced over their placeholders. + h = nn.functional.embedding(input_ids, core.w["embed"]) + eager_on_graph(h, inner=self._splice) + dt = h.dtype + dev = h.device + + # Scratch shared by every layer that uses it; each break rewrites its whole output + # and both packed buffers are fully rewritten by the projection copies. qkv = torch.empty([num_tokens, core.qkv_width], dtype=dt, device=dev) index_scratch = torch.empty([num_tokens, core.index_scratch_width], dtype=dt, device=dev) attn_out = torch.empty([num_tokens, core.q_width], dtype=dt, device=dev) + gdn_out = torch.empty( + [1, num_tokens, core.num_v_heads, core.head_v_dim], dtype=dt, device=dev + ) + ple_out = torch.empty([num_tokens, core.stream], dtype=dt, device=dev) + # The CuteDSL routed output, shared by all layers: each layer's unpermute rewrites + # every row, and the post-inject consumes it before the next layer's MoE starts. + routed_out = ( + torch.empty([num_tokens, core.hidden], dtype=dt, device=dev) + if core.trunk_moe_backend == "CUTEDSL" + else None + ) # The embedding repeated across the hc_count residual streams. hc_count = core.hc_count hc = torch.cat([h] * hc_count, -1) for i, L in enumerate(self._layers): if i == core.ple_layer: - hc = torch.add(hc, self._ple_block(hc, input_ids, ple_state)) + eager_on_graph(hc, ple_out, inner=self._ple_block) + hc = torch.add(hc, ple_out) 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 + mixed, i, L, qkv, index_scratch, attn_out, pos, pos3 ) else: - block_out = self._gdn_block(mixed, i, L, gdn_state) + block_out = self._gdn_block(mixed, i, L, gdn_out) 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) + hc = qwen4_exp_post_inject(hc, self._moe_block(mixed, L, dt, routed_out), 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: + if not core.mtp_enabled: return final # `lm_head` reads the `[T, hidden]` mix; the MTP module reads the - # `[T, hc_count * hidden]` stream it was collapsed from. - return final, hc + # `[T, hc_count * hidden]` stream it was collapsed from. They leave packed in one + # tensor, because a breakable prefill graph's body returns one tensor; the shell + # takes the two column views (`split_spec_output`). + return torch.cat([final, hc], -1) # ------------------------------------------------------------- the draft --- @@ -2497,6 +3491,36 @@ def _mtp_moe_block(self, x, dt) -> torch.Tensor: gate_logits = cublas_mm(x, M.sh_gate) return fused_sigmoid_gate_mul_add(routed, gate_logits, shared) + def _capture_selection(self, md, topk_indices, visible_blocks) -> None: + """IndexShare: stash, per request, the selection of the row the draft loop continues + from, with that row's complete-group count, in the draft layer's row of the shared + buffers. + + That row is the one the worker gathers for the next draft step: a context + request's last row of this chunk, and a generation request's last accepted row + (`num_accepted - 1` past its first). Capturing here rather than at the first draft + decode step keeps the indexer at one call per draft loop for a batch with context + rows too.""" + batch = md.num_seqs + assert md.qsa_mtp_num_accepted is not None, ( + "IndexShare needs the worker's accepted counts to pick each request's capture " + "row; the draft loop did not supply them" + ) + seq_lens = md.seq_lens_cuda[:batch].to(torch.int64) + starts = torch.cumsum(seq_lens, 0) - seq_lens + accepted = md.qsa_mtp_num_accepted[:batch].to(torch.int64) + offsets = torch.minimum(torch.clamp(accepted - 1, min=0), seq_lens - 1) + offsets[: md.num_contexts] = seq_lens[: md.num_contexts] - 1 + rows = starts + offsets + slot = self.core._draft_bound.share_slot + md.qsa_shared_topk_indices[slot, :batch].copy_(torch.index_select(topk_indices, 0, rows)) + md.qsa_shared_topk_visible_blocks[slot, :batch].copy_( + torch.index_select(visible_blocks, 0, rows) + ) + # Host-side readiness, recorded in program order: under capture this runs while the + # graph is recorded, so a graph holds the reuse only after the capture it reads. + md.qsa_shared_topk_captured_slots.add(slot) + def mtp_forward( self, embed_tokens: torch.Tensor, @@ -2567,6 +3591,43 @@ def mtp_forward( "window is not implemented here" ) + # The draft layer's QSA planes and the page table that addresses them. With a + # separate draft manager the engine's QSA metadata swaps `qsa_block_table` for the + # table it decodes from the draft manager's offsets (`prepare_for_draft_forward`); + # a table still decoding the trunk's offsets would read and write another pool's + # pages. + draft = core._draft_bound + if draft is None or cm is not draft.cm: + core._bind_draft_cache_manager(md) + draft = core._draft_bound + assert cm is core._bound.cm or md.qsa_block_table is getattr( + md, "qsa_draft_block_table", None + ), ( + "the draft step is bound to a separate draft KV cache manager but " + "qsa_block_table is not that manager's table: the engine's QSA metadata did " + "not swap its page table for the draft forward" + ) + # The `qsa_*` row mapping is sized from the metadata's token count, which the + # worker rewrites together with it between draft steps. + assert rows == md.num_tokens, ( + f"the draft step carries {rows} rows for a batch its metadata describes as " + f"{md.num_tokens} tokens; the QSA row mapping covers the second" + ) + # The trunk's dense/sparse decision from the same host buffer: the worker advances + # only the device KV lengths through the draft loop, so every draft step of one + # engine forward takes the branch its trunk took, in the CUDA-graph family the + # runner chose for that trunk. + batch = md.num_seqs + sparse = batch > 0 and int(md.kv_lens_runtime[:batch].max()) > core.dense_threshold + # IndexShare: the only draft step with context rows is the loop's first, its + # draft-extend step. It always runs the indexer and, with sharing on, captures the + # selection the draft decode steps reuse. + share = core.mtp_index_share and md.qsa_in_mtp_draft_loop + assert not (share and md.qsa_indexer_skip_topk), ( + "the worker asked a draft step with context rows to reuse a selection; only " + "the draft decode steps after it reuse" + ) + # Copy into a fresh contiguous `[3, N]`: a sliced view keeps the source row # stride and the op rejects it under `use_mrope`. assert position_ids.shape[0] == _QSA_POSITION_AXES, ( @@ -2579,6 +3640,10 @@ def mtp_forward( device=position_ids.device, ) pos3.copy_(torch.reshape(position_ids, [_QSA_POSITION_AXES, -1])) + # The trunk's two other views, from the copy: `pos` (coordinate 0 alone, which the + # index rope reads) and `coords` (the `[N, 3]` transpose the position plane stores). + pos = torch.reshape(torch.split(pos3, [1, _QSA_POSITION_AXES - 1], 0)[0], [-1]) + coords = torch.transpose(pos3, 0, 1) # ---- the front end -------------------------------------------------- e = nn.functional.embedding(input_ids, embed_tokens) @@ -2637,19 +3702,139 @@ def mtp_forward( mrope_section1=core.mrope_section1, mrope_section2=core.mrope_section2, ) - # Dense attention; see `MTPLayer` for why the module's QSA indexer is unused. - attn_out = thop_attention( - q=qkv, - output=attn_out, - local_layer_idx=layer_idx, - predicted_tokens_per_seq=p, - num_heads=core.heads_q, - num_kv_heads=core.heads_kv, - head_size=core.head_dim, - quant_mode=core.quant_mode, - **step, - **_CALL_CONSTANTS, + # The QSA index branch, as in the trunk's `_attention_block`: the raw key stored + # before anything normalizes or rotates it, the query's own norm and rotation, + # then the compression of every group this step completes. Maintained on every + # draft step at both ends of the threshold, so the draft plane is complete when a + # request crosses it. + req_idx = md.qsa_req_idx_per_token[:rows] + logical = md.qsa_logical_positions[:rows] + block_table = md.qsa_block_table + tokens_per_block = md.tokens_per_block + index_cache = draft.index_cache + qk_index = cublas_mm(mixed, M.idx) + token_k = torch.split( + torch.reshape(qk_index, [rows, core.index_heads + core.index_kv_heads, -1]), + [core.index_heads, core.index_kv_heads], + 1, + )[1] + qsa_paged_index_store( + token_k, + coords, + index_cache, + draft.position_cache, + req_idx, + logical, + block_table, + tokens_per_block, ) + index_scratch = torch.empty([rows, core.index_scratch_width], dtype=dt, device=dev) + index_scratch[:, : core.index_width].copy_(qk_index) + fused_qk_norm_rope( + index_scratch, + num_heads_q=core.index_heads, + num_heads_k=core.index_kv_heads, + num_heads_v=core.index_kv_heads, + head_dim=core.index_head_dim, + rotary_dim=core.rotary_dim, + eps=core.eps, + q_weight=M.idx_qn, + k_weight=M.idx_kn, + base=core.mtp_theta, + is_neox=True, + position_ids=pos3 if _QSA_INDEX_USE_MROPE else pos, + is_qk_norm=True, + use_gemma=True, + use_mrope=_QSA_INDEX_USE_MROPE, + mrope_section1=(_QSA_MROPE_SECTION or (0, 0, 0))[1], + mrope_section2=(_QSA_MROPE_SECTION or (0, 0, 0))[2], + ) + q_index = torch.split( + torch.reshape( + index_scratch, + [rows, core.index_heads + 2 * core.index_kv_heads, core.index_head_dim], + ), + [core.index_heads, core.index_kv_heads, core.index_kv_heads], + 1, + )[0] + qsa_prefill_compress( + logical_positions=logical, + request_indices=req_idx, + block_table=block_table, + index_cache=index_cache, + position_cache=draft.position_cache, + k_norm_weight=M.idx_kn, + cos_sin=core._mtp_cos_sin, + eps=core.eps, + tokens_per_block=tokens_per_block, + compress_ratio=core.compress_ratio, + mrope_section=_QSA_MROPE_SECTION, + ) + + if sparse: + kv = draft.kv + # The append the dense op would have owned, then selection over the draft + # layer's own compressed keys and the sparse GQA over its pool. + _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block) + scores = qsa_paged_index_scores( + q=q_index, + index_cache=index_cache, + block_table=block_table, + query_positions=logical, + request_indices=req_idx, + tokens_per_block=tokens_per_block, + compress_ratio=core.compress_ratio, + only_visible_blocks=True, + context_rows=True, + ) + topk_indices = md.qsa_topk_indices[:rows] + visible_blocks = md.qsa_visible_blocks[:rows] + _select( + core, + scores, + md.qsa_topk_row_starts[:rows], + visible_blocks, + topk_indices, + _live_score_columns(core, md, batch), + ) + if share: + self._capture_selection(md, topk_indices, visible_blocks) + selected = qsa_expand_block_indices( + topk_indices, + logical, + md.qsa_sequence_lengths[:rows], + compress_ratio=core.compress_ratio, + token_topk=core.token_topk, + ) + attn_out = torch.reshape( + _sparse_gqa( + core, + kv, + q=q_dst, + block_table=block_table, + selected_tokens=selected, + request_indices=req_idx, + tokens_per_block=tokens_per_block, + query_positions=logical, + compress_ratio=core.compress_ratio, + ), + [rows, core.q_width], + ) + else: + attn_out = thop_attention( + q=qkv, + output=attn_out, + local_layer_idx=layer_idx, + predicted_tokens_per_seq=p, + num_heads=core.heads_q, + num_kv_heads=core.heads_kv, + head_size=core.head_dim, + quant_mode=core.quant_mode, + kv_scale_orig_quant=core.kv_scale_orig_quant, + kv_scale_quant_orig=core.kv_scale_quant_orig, + **step, + **_CALL_CONSTANTS, + ) fused_sigmoid_mul_inplace(attn_out, gate) x = qwen4_exp_post_inject(x, cublas_mm(attn_out, M.o), inject) @@ -2668,6 +3853,11 @@ def __init__(self, core: Qwen4ExpTp1ModelingV2) -> None: self.core = core # See PrefillTarget.__init__. self._contract_pending = step_contract_enabled() + self._kv_probe = ( + _KvRangeProbe(core.attn_layers[-1]) + if self._contract_pending and core.quant_mode == QM_FP8_KV_CACHE + else None + ) self._layers = _layer_views(core) self._stack = _stack_views(core) @@ -2737,49 +3927,117 @@ def _grouped_norm(self, x: torch.Tensor, weights: tuple) -> torch.Tensor: [flashinfer_gemma_rmsnorm(part, w, core.eps) for w, part in zip(weights, parts)], 1 ) - def _moe_block(self, x: torch.Tensor, L: _LayerWeights, dt: torch.dtype) -> torch.Tensor: + def _moe_block( + self, x: torch.Tensor, L: _LayerWeights, dt: torch.dtype, routed_out: torch.Tensor | None + ) -> 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] + if core.trunk_moe_backend == "CUTEDSL": + # fp32 weights: the dtype the unpermute scales by. + topk_ids, topk_w = renorm_moe_routing_op( + logits, topk=core.topk, output_dtype=torch.float32 + ) + xq, xsf = fp4_quantize(x, L.moe_gscale, _SF_VEC, False, _SF_LINEAR) + # The sort pads each expert's (token, expert) pairs to whole tiles and emits + # the tile tables both GEMMs walk. + tile_expert, tile_limit, expanded_to_permuted, permuted_to_expanded, _, live_tiles = ( + moe_sort( + token_selected_experts=topk_ids, + token_final_scales=topk_w, + num_experts=core.num_experts, + top_k=core.topk, + local_expert_offset=0, + local_num_experts=core.num_experts, + tile_tokens_dim=_CUTEDSL_TILE_SIZE, + ) + ) + # FC1 gathers its rows from the unpermuted activation, applies SwiGLU and + # requantizes with the layer's one FC2 input scale. `fp4_quantize` returns the + # linear scales flat; the GEMM reads them as [T, hidden / 16]. + h, h_sf = cute_dsl_nvfp4_gather_grouped_gemm_act_fusion( + input=xq.view(torch.float4_e2m1fn_x2), + weight=L.fc1_w, + input_scale=xsf.view(x.shape[0], core.hidden // _SF_VEC), + weight_scale=L.fc1_s, + alpha=L.alpha1, + tile_idx_to_group_idx=tile_expert, + tile_idx_to_mn_limit=tile_limit, + permuted_idx_to_expanded_idx=permuted_to_expanded, + num_non_exiting_tiles=live_tiles, + global_sf=L.global_sf, + num_experts=core.num_experts, + top_k=core.topk, + num_local_experts=core.num_experts, + local_expert_offset=0, + tile_size=_CUTEDSL_TILE_SIZE, + scaling_vector_size=_SF_VEC, + activation_type=int(ActivationType.Swiglu), + ) + # FC2 writes one bf16 row per permuted row; the unpermute writes every row of + # `routed_out` as the routing-weighted sum of its token's top-k rows. + y = cute_dsl_nvfp4_grouped_gemm( + input=h, + weight=L.fc2_w, + input_scale=h_sf, + weight_scale=L.fc2_s, + alpha=L.alpha2, + tile_idx_to_group_idx=tile_expert, + num_non_exiting_tiles=live_tiles, + num_experts=core.num_experts, + top_k=core.topk, + num_local_experts=core.num_experts, + local_expert_offset=0, + tile_size=_CUTEDSL_TILE_SIZE, + output_dtype=torch.bfloat16, + scaling_vector_size=_SF_VEC, + ) + routed = moe_unpermute_inplace( + permuted_input=y, + output=routed_out, + expanded_idx_to_permuted_idx=expanded_to_permuted, + topk_scales=topk_w, + ) + else: + 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): + def _attention_block(self, x, i, L, qkv, index_scratch, attn_out, qsa, step, stage): """One gated GQA layer plus its QSA index branch for a generation-only batch; see `PrefillTarget._attention_block`.""" core = self.core @@ -2835,6 +4093,8 @@ def _attention_block(self, x, i, L, qkv, index_scratch, attn_out, qsa, step): mrope_section1=core.mrope_section1, mrope_section2=core.mrope_section2, ) + if self._kv_probe is not None and not self._kv_probe.observe(i, k_dst, v_dst): + self._kv_probe = None # ---- the QSA index branch (see PrefillTarget) ------------------------ # The raw key is stored before any norm/rotation; the compression step @@ -2846,6 +4106,8 @@ def _attention_block(self, x, i, L, qkv, index_scratch, attn_out, qsa, step): [core.index_heads, core.index_kv_heads], 1, )[1] + if stage: + core.spec_rollback.stage_index(index_cache, core._bound.position_cache) qsa_paged_index_store( token_k, coords, @@ -2901,12 +4163,10 @@ def _attention_block(self, x, i, L, qkv, index_scratch, attn_out, qsa, step): ) if sparse: - k_cache, v_cache = core._bound.kv[i] + kv = core._bound.kv[i] # Append the fused-norm+rope'd K/V; 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 - ) + # token from the pool. Over an fp8 or NVFP4 pool as in the prefill target. + _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block) scores = qsa_paged_index_scores( q=q_index, index_cache=index_cache, @@ -2918,10 +4178,7 @@ def _attention_block(self, x, i, L, qkv, index_scratch, attn_out, qsa, step): 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 - ) + _select(core, scores, topk_row_starts, visible_blocks, topk_indices) selected = qsa_expand_block_indices( topk_indices, logical, @@ -2930,15 +4187,14 @@ def _attention_block(self, x, i, L, qkv, index_scratch, attn_out, qsa, step): token_topk=core.token_topk, ) out = torch.reshape( - qsa_paged_sparse_gqa( + _sparse_gqa( + core, + kv, q=q_dst, - k_cache=k_cache, - v_cache=v_cache, block_table=block_table, selected_tokens=selected, request_indices=req_idx, tokens_per_block=tokens_per_block, - softmax_scale=core.softmax_scale, # The prefill pair bounds a context row's scan; none here. query_positions=None, compress_ratio=None, @@ -2954,13 +4210,17 @@ def _attention_block(self, x, i, L, qkv, index_scratch, attn_out, qsa, step): num_kv_heads=core.heads_kv, head_size=core.head_dim, quant_mode=core.quant_mode, + kv_scale_orig_quant=core.kv_scale_orig_quant, + kv_scale_quant_orig=core.kv_scale_quant_orig, **step, **_CALL_CONSTANTS, ) fused_sigmoid_mul_inplace(out, gate) return cublas_mm(out, L.o) - def _gdn_verify(self, mq_d, a_d, b_d, ssm_states, i, L, si_d, int_idx, n, p) -> torch.Tensor: + def _gdn_verify( + self, mq_d, a_d, b_d, ssm_states, i, L, si_d, int_idx, n, p, replay + ) -> torch.Tensor: """Recurrent half of a speculative verify step; see `PrefillTarget._gdn_verify`.""" core = self.core @@ -2982,6 +4242,8 @@ def _gdn_verify(self, mq_d, a_d, b_d, ssm_states, i, L, si_d, int_idx, n, p) -> ) # Back to token-major (a view). mq_out = torch.reshape(torch.transpose(conv_out, 1, 2), [n * p, core.conv_dim]) + if replay is not None: + return _replay_verify(core, mq_out, a_d, b_d, ssm_states, i, L, si_d, n, p, replay) q, k, v = pack_gdn_decode_qkv( mq_out, core.num_k_heads, core.head_k_dim, core.num_v_heads, core.head_v_dim ) @@ -3025,7 +4287,7 @@ 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 + si_all, qsl_long, is_verify, num_decodes, gen_p, int_idx, replay = gdn_state num_tokens = x.shape[0] mixed_qkv = cublas_mm(x, L.in_qkv) @@ -3036,7 +4298,7 @@ def _gdn_block(self, x, i, L, gdn_state) -> torch.Tensor: 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 + mixed_qkv, a, b, ssm_states, i, L, si_all, int_idx, num_decodes, gen_p, replay ) else: causal_conv1d_update( @@ -3049,25 +4311,55 @@ def _gdn_block(self, x, i, L, gdn_state) -> torch.Tensor: conv_state_indices=si_all, pad_slot_id=_PAD_SLOT_ID, ) - # q/k/v stay views over the conv output (native strides, no copy). - qv, kv, vv = torch.split(mixed_qkv, [core.key_dim, core.key_dim, core.value_dim], 1) - 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, - ) + if ssm_states.dtype == torch.bfloat16: + # Over a bf16 pool a call with `cu_seqlens` reaches FlashInfer's CuTe-DSL + # decode, whose kernel changes with the batch size and so changes a row's + # last bits; without `cu_seqlens` the entry runs its Triton kernel, one program + # per (row, head). That form finds row n's tokens at `n * q.stride(1)`, so q/k/v + # must be split from the whole contiguous `[T, 10240]` block reshaped to + # `[T, 1, 80, 128]` (K and V heads are both 128 wide), where the size-1 token + # axis carries the row stride; a/b are contiguous `[T, 48]`. + qkv = torch.reshape(mixed_qkv, [num_tokens, 1, -1, core.head_k_dim]) + q, k, v = torch.split( + qkv, [core.num_k_heads, core.num_k_heads, core.num_v_heads], 2 + ) + out = fused_sigmoid_gating_delta_rule_update( + A_log=L.a_log, + a=torch.reshape(a, [num_tokens, 1, core.num_v_heads]), + dt_bias=L.dt_bias, + softplus_beta=_SOFTPLUS_BETA, + softplus_threshold=_SOFTPLUS_THRESHOLD, + q=q, + k=k, + v=v, + b=torch.reshape(b, [num_tokens, 1, core.num_v_heads]), + initial_state_source=ssm_states, + initial_state_indices=si_all, + scale=None, + use_qk_l2norm_in_kernel=True, + cu_seqlens=None, + output=None, + ) + else: + # 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: @@ -3084,12 +4376,14 @@ def _gdn_epilogue(self, out, z, L, num_tokens) -> torch.Tensor: ) return cublas_mm(torch.reshape(o, [num_tokens, core.value_dim]), L.out) - def _ple_block(self, hc, input_ids, ple_state) -> torch.Tensor: + def _ple_block(self, hc, input_ids, ple_state, stage) -> 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 + if stage: + core.spec_rollback.stage_ple_pools(ngram_pool, conv_states) window = ngram_token_history( input_ids, ngram_pool, si_all, carry_all, qsl, core.eos_token_id, core.ngram_size ) @@ -3115,9 +4409,12 @@ def _ple_block(self, hc, input_ids, ple_state) -> torch.Tensor: gated = gated_ngram_injection( key_normed, query_normed, value, core.hc_count, _PLE_CLAMP_MIN ) + conv_in = self._grouped_norm(gated, S.ple_nc) + if stage: + core.spec_rollback.stage_ple_rows(window, conv_in) # The short conv reads channel-major; a bare transpose is the certified form. convolved = dilated_short_conv_paged( - x=torch.transpose(self._grouped_norm(gated, S.ple_nc), 0, 1), + x=torch.transpose(conv_in, 0, 1), weight=S.ple_conv_w, bias=None, conv_states=conv_states, @@ -3209,21 +4506,22 @@ def forward( "per-step conv snapshots and the recurrent rollback buffer are " "sized from the second and written from the first" ) - assert 1 <= gen_p <= 4, ( - f"predicted_tokens_per_seq {gen_p} is outside the 1..4 range " - f"attention/thop_attention certifies" - ) + check_verify_width(gen_p) + check_replay_width(core._bound.replay, gen_p) is_verify = spec_metadata is not None and batch > 0 and gen_p > 1 step["predicted_tokens_per_seq"] = gen_p si_all = mm.state_indices[:batch] qsl_long = mm.query_start_loc_long int_idx = None + replay = None if is_verify: # The runtime's arange over the slot pool (what `update_mamba_states` reads). assert core._bound.int_idx is not None, "the cache-manager bind has not run" int_idx = core._bound.int_idx[:batch] - gdn_state = (si_all, qsl_long, is_verify, batch, gen_p, int_idx) + if core._bound.replay is not None: + replay = _replay_work(core, mm, batch) + gdn_state = (si_all, qsl_long, is_verify, batch, gen_p, int_idx, replay) # Every generation sequence carries its n-gram history; the engine's # `has_initial_states` covers the context prefix only (empty here). @@ -3257,23 +4555,49 @@ def forward( sparse, ) + # Staging for the speculative rollback, as on the prefill target (see + # PrefillTarget.forward); every row here is a generation row. + stage = is_verify + rollback = core.spec_rollback + if rollback is not None: + rollback.discard() + if stage: + rollback.begin( + num_contexts=0, + prefix_tokens=0, + num_gens=batch, + rows_per_gen=gen_p, + slots=si_all, + request_indices=md.qsa_req_idx_per_token[:num_tokens], + logical_positions=md.qsa_logical_positions[:num_tokens], + block_table=md.qsa_block_table, + tokens_per_block=md.tokens_per_block, + ) + # Scratch shared by all attention layers (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 CuteDSL routed output, shared by all layers: each layer's unpermute rewrites + # every row, and the post-inject consumes it before the next layer's MoE starts. + routed_out = ( + torch.empty([num_tokens, core.hidden], dtype=dt, device=dev) + if core.trunk_moe_backend == "CUTEDSL" + else None + ) # The embedding repeated across the hc_count residual streams. hc_count = core.hc_count hc = torch.cat([h] * hc_count, -1) for i, L in enumerate(self._layers): if i == core.ple_layer: - hc = torch.add(hc, self._ple_block(hc, input_ids, ple_state)) + hc = torch.add(hc, self._ple_block(hc, input_ids, ple_state, stage)) 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 + mixed, i, L, qkv, index_scratch, attn_out, qsa, step, stage ) else: block_out = self._gdn_block(mixed, i, L, gdn_state) @@ -3281,7 +4605,7 @@ def forward( 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) + hc = qwen4_exp_post_inject(hc, self._moe_block(mixed, L, dt, routed_out), inject) # The stack-final mixer with no injection weight is the final norm. S = self._stack final, _ = qwen4_exp_gated_residual( @@ -3330,6 +4654,26 @@ def _mtp_moe_block(self, x, dt) -> torch.Tensor: gate_logits = cublas_mm(x, M.sh_gate) return fused_sigmoid_gate_mul_add(routed, gate_logits, shared) + def _capture_selection(self, md, topk_indices, visible_blocks, batch, p) -> None: + """IndexShare: stash, per request, the selection of its last accepted draft-extend + row with that row's complete-group count, in the draft layer's row of the shared + buffers. Every request carries `p` rows here, and the worker continues the draft + loop from row `num_accepted - 1` of each (the golden token counts as accepted).""" + assert md.qsa_mtp_num_accepted is not None, ( + "IndexShare needs the worker's accepted counts to pick each request's capture " + "row; the draft loop did not supply them" + ) + accepted = md.qsa_mtp_num_accepted[:batch].to(torch.int64) + rows = torch.arange(batch, dtype=torch.int64, device=topk_indices.device) * p + rows = rows + torch.clamp(accepted - 1, min=0, max=p - 1) + slot = self.core._draft_bound.share_slot + md.qsa_shared_topk_indices[slot, :batch].copy_(torch.index_select(topk_indices, 0, rows)) + md.qsa_shared_topk_visible_blocks[slot, :batch].copy_( + torch.index_select(visible_blocks, 0, rows) + ) + # Host-side readiness; see PrefillTarget._capture_selection. + md.qsa_shared_topk_captured_slots.add(slot) + def mtp_forward( self, embed_tokens: torch.Tensor, @@ -3396,6 +4740,44 @@ def mtp_forward( "window is not implemented here" ) + # The draft layer's QSA planes and the page table that addresses them. With a + # separate draft manager the engine's QSA metadata swaps `qsa_block_table` for the + # table it decodes from the draft manager's offsets (`prepare_for_draft_forward`); + # a table still decoding the trunk's offsets would read and write another pool's + # pages. + draft = core._draft_bound + if draft is None or cm is not draft.cm: + core._bind_draft_cache_manager(md) + draft = core._draft_bound + assert cm is core._bound.cm or md.qsa_block_table is getattr( + md, "qsa_draft_block_table", None + ), ( + "the draft step is bound to a separate draft KV cache manager but " + "qsa_block_table is not that manager's table: the engine's QSA metadata did " + "not swap its page table for the draft forward" + ) + # The `qsa_*` row mapping is sized from the metadata's token count, which the + # worker rewrites together with it between draft steps. + assert rows == md.num_tokens, ( + f"the draft step carries {rows} rows for a batch its metadata describes as " + f"{md.num_tokens} tokens; the QSA row mapping covers the second" + ) + # The trunk's dense/sparse decision from the same host buffer: the worker advances + # only the device KV lengths through the draft loop, so every draft step of one + # engine forward takes the branch its trunk took, in the CUDA-graph family the + # runner chose for that trunk. + sparse = gen_seqs > 0 and int(md.kv_lens_runtime[:gen_seqs].max()) > core.dense_threshold + # IndexShare: the loop's draft-extend step (`p` rows per request) runs the indexer + # and, with sharing on, captures each request's selection; every draft decode step + # after it (one row per request) reuses that capture instead of running the indexer. + share = core.mtp_index_share and md.qsa_in_mtp_draft_loop + reuse = share and md.qsa_indexer_skip_topk + assert not share or reuse == (p == 1), ( + f"the worker's reuse flag ({md.qsa_indexer_skip_topk}) disagrees with the " + f"step's shape ({p} rows per request): only the draft decode steps after the " + "draft-extend step reuse its selection" + ) + # Copy into a fresh contiguous `[3, N]`: a sliced view keeps the source row # stride and the op rejects it under `use_mrope`. assert position_ids.shape[0] == _QSA_POSITION_AXES, ( @@ -3408,6 +4790,10 @@ def mtp_forward( device=position_ids.device, ) pos3.copy_(torch.reshape(position_ids, [_QSA_POSITION_AXES, -1])) + # The trunk's two other views, from the copy: `pos` (coordinate 0 alone, which the + # index rope reads) and `coords` (the `[N, 3]` transpose the position plane stores). + pos = torch.reshape(torch.split(pos3, [1, _QSA_POSITION_AXES - 1], 0)[0], [-1]) + coords = torch.transpose(pos3, 0, 1) # ---- the front end -------------------------------------------------- e = nn.functional.embedding(input_ids, embed_tokens) @@ -3466,19 +4852,150 @@ def mtp_forward( mrope_section1=core.mrope_section1, mrope_section2=core.mrope_section2, ) - # Dense attention; see `MTPLayer` for why the module's QSA indexer is unused. - attn_out = thop_attention( - q=qkv, - output=attn_out, - local_layer_idx=layer_idx, - predicted_tokens_per_seq=p, - num_heads=core.heads_q, - num_kv_heads=core.heads_kv, - head_size=core.head_dim, - quant_mode=core.quant_mode, - **step, - **_CALL_CONSTANTS, - ) + req_idx = md.qsa_req_idx_per_token[:rows] + logical = md.qsa_logical_positions[:rows] + block_table = md.qsa_block_table + tokens_per_block = md.tokens_per_block + index_cache = draft.index_cache + # The QSA index branch, as in the trunk's `_attention_block` (see the prefill + # target's draft step). Under IndexShare a draft decode step runs none of it: what + # it would write (its row's raw key and coordinates, and the anchor of a group that + # row completes) sits at or past the newly committed length, and nothing reads it + # before the next engine step's draft-extend rewrites those positions, storing + # before it compresses and compressing before it scores. + if not reuse: + qk_index = cublas_mm(mixed, M.idx) + token_k = torch.split( + torch.reshape(qk_index, [rows, core.index_heads + core.index_kv_heads, -1]), + [core.index_heads, core.index_kv_heads], + 1, + )[1] + qsa_paged_index_store( + token_k, + coords, + index_cache, + draft.position_cache, + req_idx, + logical, + block_table, + tokens_per_block, + ) + index_scratch = torch.empty([rows, core.index_scratch_width], dtype=dt, device=dev) + index_scratch[:, : core.index_width].copy_(qk_index) + fused_qk_norm_rope( + index_scratch, + num_heads_q=core.index_heads, + num_heads_k=core.index_kv_heads, + num_heads_v=core.index_kv_heads, + head_dim=core.index_head_dim, + rotary_dim=core.rotary_dim, + eps=core.eps, + q_weight=M.idx_qn, + k_weight=M.idx_kn, + base=core.mtp_theta, + is_neox=True, + position_ids=pos3 if _QSA_INDEX_USE_MROPE else pos, + is_qk_norm=True, + use_gemma=True, + use_mrope=_QSA_INDEX_USE_MROPE, + mrope_section1=(_QSA_MROPE_SECTION or (0, 0, 0))[1], + mrope_section2=(_QSA_MROPE_SECTION or (0, 0, 0))[2], + ) + q_index = torch.split( + torch.reshape( + index_scratch, + [rows, core.index_heads + 2 * core.index_kv_heads, core.index_head_dim], + ), + [core.index_heads, core.index_kv_heads, core.index_kv_heads], + 1, + )[0] + qsa_prefill_compress( + logical_positions=logical, + request_indices=req_idx, + block_table=block_table, + index_cache=index_cache, + position_cache=draft.position_cache, + k_norm_weight=M.idx_kn, + cos_sin=core._mtp_cos_sin, + eps=core.eps, + tokens_per_block=tokens_per_block, + compress_ratio=core.compress_ratio, + mrope_section=_QSA_MROPE_SECTION, + ) + + if sparse: + kv = draft.kv + _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block) + topk_indices = md.qsa_topk_indices[:rows] + visible_blocks = md.qsa_visible_blocks[:rows] + if reuse: + # The capture this loop's draft-extend step made, with the groups completed + # since appended; the expansion then adds the incomplete tail, so the + # positions drafted since the capture stay visible. + slot = draft.share_slot + assert slot in md.qsa_shared_topk_captured_slots, ( + "a draft decode step would reuse a selection that no draft-extend step " + "of this loop captured" + ) + qsa_mtp_reuse_topk( + shared_indices=md.qsa_shared_topk_indices[slot], + shared_visible_blocks=md.qsa_shared_topk_visible_blocks[slot], + request_indices=req_idx, + visible_blocks=visible_blocks, + output=topk_indices, + ) + else: + scores = qsa_paged_index_scores( + q=q_index, + index_cache=index_cache, + block_table=block_table, + query_positions=logical, + request_indices=req_idx, + tokens_per_block=tokens_per_block, + compress_ratio=core.compress_ratio, + only_visible_blocks=True, + context_rows=False, + ) + _select(core, scores, md.qsa_topk_row_starts[:rows], visible_blocks, topk_indices) + if share: + self._capture_selection(md, topk_indices, visible_blocks, gen_seqs, p) + selected = qsa_expand_block_indices( + topk_indices, + logical, + md.qsa_sequence_lengths[:rows], + compress_ratio=core.compress_ratio, + token_topk=core.token_topk, + ) + attn_out = torch.reshape( + _sparse_gqa( + core, + kv, + q=q_dst, + block_table=block_table, + selected_tokens=selected, + request_indices=req_idx, + tokens_per_block=tokens_per_block, + # The prefill pair bounds a context row's scan; none here. + query_positions=None, + compress_ratio=None, + ), + [rows, core.q_width], + ) + else: + attn_out = thop_attention( + q=qkv, + output=attn_out, + local_layer_idx=layer_idx, + predicted_tokens_per_seq=p, + num_heads=core.heads_q, + num_kv_heads=core.heads_kv, + head_size=core.head_dim, + quant_mode=core.quant_mode, + kv_scale_orig_quant=core.kv_scale_orig_quant, + kv_scale_quant_orig=core.kv_scale_quant_orig, + **step, + **_CALL_CONSTANTS, + ) fused_sigmoid_mul_inplace(attn_out, gate) x = qwen4_exp_post_inject(x, cublas_mm(attn_out, M.o), inject) @@ -3555,15 +5072,48 @@ def _vision_views(core: Qwen4ExpTp1ModelingV2) -> _VisionWeights: ) +def _frame_calls( + grid: list[tuple[int, int, int]], max_rows: int +) -> list[list[tuple[int, int, int]]]: + """Cut `grid` (one `(frames, height, width)` row per image or video) into tower calls + of whole frames, in order, each at most `max_rows` patch rows. A call is its grid + rows: the frames of one item that fall in it are one `(frames, height, width)` row.""" + calls: list[list[tuple[int, int, int]]] = [[]] + room = max_rows + for frames, height, width in grid: + rows = height * width + assert rows <= max_rows, ( + f"one frame is {height}x{width} = {rows} patch rows; the vision ops are " + f"certified up to {max_rows} rows in one call" + ) + while frames: + take = min(frames, room // rows) + if take == 0: + calls.append([]) + room = max_rows + continue + calls[-1].append((take, height, width)) + room -= take * rows + frames -= take + return calls + + class VisionTower: - """One request's images through the 27-block ViT and its patch merger. + """One request's images or videos through the 27-block ViT and its patch merger. Not a target: runs only for a context request carrying media, from the shell's - splice. Attention is bidirectional per frame, expressed to `thop_attention` as - one context "sequence" per frame; running the rows as one sequence is 315x-582x - outside tolerance. The position-embedding interpolation and 2-D rope must both - emit in spatial-merge order (what the merger's 4-to-1 reshape consumes); a - raster-order emitter gives a correctly shaped, silently wrong tower.""" + splice. Attention is bidirectional per frame (a video's frame group is a frame), + expressed to `thop_attention` as one context "sequence" per frame; running the + rows as one sequence is 315x-582x outside tolerance. The position-embedding + interpolation and 2-D rope must both emit in spatial-merge order (what the + merger's 4-to-1 reshape consumes); a raster-order emitter gives a correctly + shaped, silently wrong tower. + + The frames run in calls of at most `_VISION_CALL_MAX_PATCHES` rows, cut only between + frames: attention is per frame, the position embedding and rope are per grid row + and repeat per frame, the norms and MLPs are row-wise, and a frame's `h * w` is a + multiple of `merge**2`, so a cut keeps every merger block whole. One frame past the + bound cannot be served.""" def __init__(self, core: Qwen4ExpTp1ModelingV2) -> None: self.core = core @@ -3575,10 +5125,11 @@ def __init__(self, core: Qwen4ExpTp1ModelingV2) -> None: def __call__(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor: """`pixel_values` is `[sum(t*h*w), in_channels*kt*kh*kw]`, one flat patch per row, channel-slowest. `grid_thw` is one `(frames, height, width)` row per - image and stays on the host (no synchronization). Returns - `[sum(t*h*w) / merge**2, out_hidden_size]`, one row per placeholder.""" + image or video and stays on the host (no synchronization). Returns + `[sum(t*h*w) / merge**2, out_hidden_size]`, one row per placeholder, in grid + row order.""" core = self.core - dt, dev = pixel_values.dtype, pixel_values.device + dt = pixel_values.dtype assert dt == core.w["embed"].dtype, ( f"the tower runs at the activation dtype {core.w['embed'].dtype}; the " f"runtime handed pixel_values as {dt}. The input processor's `dtype` " @@ -3588,29 +5139,53 @@ def __call__(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> torch. f"pixel_values must be [patches, {core.v_pixel_width}] flat patch rows; " f"got {tuple(pixel_values.shape)}" ) - assert pixel_values.is_cuda, "pixel_values must reach the tower on the device" - # ---- the frame segmentation, host-side ------------------------------- - # One attention segment per frame: a video's frames do not attend to each other. - grid = grid_thw.tolist() - seg_lens: list[int] = [] - for row in grid: + grid: list[tuple[int, int, int]] = [] + for row in grid_thw.tolist(): assert len(row) == 3, f"grid_thw rows are (t, h, w); got {row}" - frames, height, width = int(row[0]), int(row[1]), int(row[2]) + frames, height, width = (int(v) for v in row) assert frames >= 1 and height >= 1 and width >= 1, row assert height % core.v_merge == 0 and width % core.v_merge == 0, ( f"grid sides must divide by spatial_merge_size {core.v_merge}; got " f"height={height}, width={width}" ) + grid.append((frames, height, width)) + assert grid, "grid_thw describes no image or video" + described = sum(t * h * w for t, h, w in grid) + assert described == pixel_values.shape[0], ( + f"grid_thw describes {described} patch rows and pixel_values carries " + f"{pixel_values.shape[0]}" + ) + + outputs = [] + start = 0 + for call in _frame_calls(grid, _VISION_CALL_MAX_PATCHES): + rows = sum(t * h * w for t, h, w in call) + outputs.append( + self._run(pixel_values[start : start + rows], call, grid_thw.new_tensor(call)) + ) + start += rows + return outputs[0] if len(outputs) == 1 else torch.cat(outputs, 0) + + def _run( + self, pixel_values: torch.Tensor, grid: list[tuple[int, int, int]], grid_thw: torch.Tensor + ) -> torch.Tensor: + """One tower call over whole frames: `grid` and `grid_thw` (the same rows, as a + host list and a host tensor) describe exactly the rows of `pixel_values`.""" + core = self.core + dt, dev = pixel_values.dtype, pixel_values.device + assert pixel_values.is_cuda, "pixel_values must reach the tower on the device" + + # ---- the frame segmentation, host-side ------------------------------- + # One attention segment per frame: a video's frames do not attend to each other. + seg_lens: list[int] = [] + for frames, height, width in grid: seg_lens.extend([height * width] * frames) num_patches = pixel_values.shape[0] - assert sum(seg_lens) == num_patches, ( - f"grid_thw describes {sum(seg_lens)} patch rows and pixel_values carries {num_patches}" - ) num_segments = len(seg_lens) max_seg = max(seg_lens) - # Batch-state tensors for the encoder-only call; they describe the images. + # Batch-state tensors for the encoder-only call; they describe the frames. host_lens = torch.tensor(seg_lens, dtype=torch.int32) host_types = torch.zeros([num_segments], dtype=torch.int32) host_total = torch.tensor([num_patches, 0], dtype=torch.int32) @@ -3681,12 +5256,26 @@ class MTPLayer: The MoE is FP8 128x128 weight blocks (the trunk's is NVFP4); re-quantizing it would change the checkpoint's meaning, visible only as lower acceptance. - Attention is dense, the one deliberate departure from the checkpoint (its QSA - indexer stays unloaded): with a separate draft KV manager the worker swaps - `attn_metadata.kv_cache_manager`, `QSAAttentionMetadata` implements no - `prepare_for_draft_forward`, so `qsa_block_table` would decode the target - manager's page offsets against the draft pool. Below `dense_seq_len_threshold` - dense is bit-for-bit the trunk's. + Attention is the trunk's QSA layer over the module's own indexer. Every draft step + maintains the layer's index plane, then runs the dense core while the trunk's host KV + lengths stay within `dense_seq_len_threshold`, and selection plus sparse GQA above it. + With a separate draft KV manager the layer's K/V, index and position pages live in + that manager's pools: the worker swaps `attn_metadata.kv_cache_manager` and the block + offsets for the draft loop, `QSAAttentionMetadata.prepare_for_draft_forward` swaps + `qsa_block_table` to the table it decodes from the draft manager's offsets, and the + core binds the draft manager's planes once per manager (`_DraftBound`). The planes need + no rollback: each engine step's first draft call rewrites every position from the + newly committed length on. + + IndexShare (`sparse_attention_config.index_share_for_mtp_iteration`) runs the draft + indexer once per MTP iteration instead of once per draft step. The draft-extend step, + the loop's first, runs it and captures per request the selection of the row the loop + continues from, into the metadata's shared buffers; every draft decode step after it + skips the index branch and rebuilds its selection from that capture with + `qsa_mtp_reuse_topk`, which appends the groups completed since; a full row gives up + its oldest ids to them. The expansion adds the incomplete tail, so the tokens drafted + since the capture stay visible. Every selection is canonical first (ascending, + padding last), so the reuse keeps the order. Not an `nn.Module`: the weights live in the core's `ParameterDict`, and a second registration would break the weight load's both-directions coverage assert.""" @@ -3758,6 +5347,170 @@ def embed_tokens(self) -> torch.Tensor: return self.core.w["embed"] +class SpeculativeRollback: + """The target's auxiliary speculative-state handler: it undoes what a verify step's + rejected draft rows wrote into the four side-cache planes the runtime's own state + promotion does not cover. + + A verify step runs `p = runtime_draft_len + 1` rows per generation request through the + trunk, and four ops write per-request state for every row without knowing how many the + sampler keeps: the PLE n-gram id history (`ops/ple/ngram_token_history`), the PLE + dilated-conv state (`ops/ssm/dilated_short_conv_paged`), each attention layer's QSA + index-K plane (`qsa_paged_index_store`, `qsa_prefill_compress`) and the shared QSA + position plane. The Gated-DeltaNet state is not here: the runtime promotes it. + + For a generation request that accepted `a` rows (golden token included, so + `1 <= a <= p`) the commit rewrites: + + * the n-gram pool row to the id window of row `a - 1` minus its oldest id; + * the conv pool row to columns `[a, a + state_len)` of `[pre-step state | conv input + of the request's p rows]`; + * every index-K and position slot written by a row of ordinal `>= a` to its + pre-verify contents. + + The targets stage during the forward (`begin`, then `stage_ple_pools` / + `stage_ple_rows` in the PLE block and `stage_index` before each attention layer's + store) and the spec worker calls `commit_speculative_states` after its draft loop, + inside the same model forward. Context rows and non-verify steps stage nothing, and a + commit with nothing staged returns. Staging and commit use metadata-sized + allocations, gathers, `torch.where` and indexed writes only, and everything staged is + allocated by the forward and consumed by that forward's commit, so a captured graph + holds both ends and replay re-reads the engine's buffers in place.""" + + def __init__(self) -> None: + self._pending: dict | None = None + + # ------------------------------------------------------------- staging --- + + def discard(self) -> None: + """Drop whatever an earlier forward staged.""" + self._pending = None + + def begin( + self, + *, + num_contexts: int, + prefix_tokens: int, + num_gens: int, + rows_per_gen: int, + slots: torch.Tensor, + request_indices: torch.Tensor, + logical_positions: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + ) -> None: + """Open a verify step's staging. The generation rows follow the `prefix_tokens` + context rows as `num_gens` runs of `rows_per_gen`, so row `s * p + t` is ordinal `t` + of generation request `s`. `slots` are those requests' recurrent slots; + `request_indices` and `logical_positions` are the generation rows' own. Each row's + QSA slot is resolved here, as `qsa_paged_index_store` resolves it.""" + n, p = num_gens, rows_per_gen + assert n > 0 and p > 1, ( + f"a verify step stages generation requests carrying drafts; got {n} " + f"requests of {p} rows" + ) + assert logical_positions.shape[0] == n * p, ( + f"{n} generation requests of {p} rows are {n * p} rows; got " + f"{logical_positions.shape[0]} positions" + ) + positions = logical_positions.to(torch.long) + pages = block_table[request_indices.to(torch.long), positions // tokens_per_block] + self._pending = dict( + num_contexts=num_contexts, + prefix=prefix_tokens, + n=n, + p=p, + slots=slots.to(torch.long), + pages=pages.to(torch.long), + within=positions % tokens_per_block, + index=[], + position=None, + ) + + def stage_ple_pools(self, ngram_pool: torch.Tensor, conv_pool: torch.Tensor) -> None: + """Before the PLE block's first write: both pools, and the generation slots' + pre-step conv rows.""" + pending = self._pending + slots = pending["slots"] + pending["ngram_pool"] = ngram_pool + pending["conv_pool"] = conv_pool + pending["conv_prev"] = conv_pool[slots] + + def stage_ple_rows(self, window: torch.Tensor, conv_input: torch.Tensor) -> None: + """After the id history and before the conv: the generation rows' id windows + (`[T, ngram_size]`, newest id last) and conv input (`[T, C]`, token-major). Copied + rather than viewed, so a long context prefix in a mixed batch is not kept alive + through the draft loop.""" + pending = self._pending + prefix, n, p = pending["prefix"], pending["n"], pending["p"] + # Row t of the window without its oldest id is the history after token t. + pending["ngram_candidates"] = torch.reshape(window[prefix:, 1:], [n, p, -1]).clone() + pending["conv_x"] = torch.reshape(conv_input[prefix:], [n, p, -1]).clone() + + def stage_index(self, index_cache: torch.Tensor, position_cache: torch.Tensor) -> None: + """Before an attention layer's index store: the generation rows' slots of its + index-K plane, and, on the step's first attention layer (the first to write the + shared plane), of the position plane.""" + pending = self._pending + pages, within = pending["pages"], pending["within"] + if pending["position"] is None: + pending["position"] = (position_cache, position_cache[pages, within]) + pending["index"].append((index_cache, index_cache[pages, within])) + + # -------------------------------------------------------------- commit --- + + def commit_speculative_states( + self, + num_accepted_tokens: torch.Tensor, + state_indices: torch.Tensor, + num_contexts: int, + ) -> None: + """Rewrite the staged planes to the accepted prefix. Called by the spec worker.""" + pending, self._pending = self._pending, None + if pending is None: + return + n, p = pending["n"], pending["p"] + assert num_contexts == pending["num_contexts"], ( + f"the worker commits with {num_contexts} context requests and the " + f"forward staged {pending['num_contexts']}" + ) + assert state_indices.shape[0] == num_contexts + n, ( + f"the worker hands {state_indices.shape[0]} slots for a batch the " + f"forward staged as {num_contexts} + {n} requests" + ) + assert pending.get("conv_x") is not None and len(pending["index"]) > 0, ( + "a verify step staged neither its PLE rows nor any attention layer; " + "the forward did not run the staging it opened" + ) + accepted = num_accepted_tokens[num_contexts : num_contexts + n].to(torch.long) + slots = state_indices[num_contexts:].to(torch.long) + rows = torch.arange(n, device=accepted.device) + ngram_pool, conv_pool = pending["ngram_pool"], pending["conv_pool"] + state_len = conv_pool.shape[-1] + ngram_pool[slots] = pending["ngram_candidates"][rows, accepted - 1] + # [n, C, state_len + p]; window k of the unfold is the state after k tokens. + history = torch.cat([pending["conv_prev"], torch.transpose(pending["conv_x"], 1, 2)], 2) + conv_pool[slots] = history.unfold(2, state_len, 1)[rows, :, accepted] + + # [n * p]: True at every ordinal the sampler rejected. + rejected = torch.reshape( + torch.arange(p, device=accepted.device).unsqueeze(0) >= accepted.unsqueeze(1), [-1] + ) + pages, within = pending["pages"], pending["within"] + for cache, old in pending["index"]: + cache[pages, within] = torch.where(rejected[:, None, None], old, cache[pages, within]) + cache, old = pending["position"] + cache[pages, within] = torch.where(rejected[:, None], old, cache[pages, within]) + + +#: Per modality, the keys under `multimodal_data[modality]`: the flat patch rows and the +#: host-side `(t, h, w)` grid. +_MEDIA_KEYS = { + "image": ("pixel_values", "image_grid_thw"), + "video": ("pixel_values_videos", "video_grid_thw"), +} + + class ModelingV2Qwen38FlashNextNvfp4Sm103Tp1( DecoderModelForCausalLM[Qwen4ExpTp1ModelingV2, PretrainedConfig] ): @@ -3844,6 +5597,10 @@ def __init__(self, model_config: ModelConfig): ) # Caches the draft->target vocab offset map (correctly None here). self.spec_worker.set_draft_model(self.draft_model) + # The PLE pools and the QSA side caches are written per verify row and nothing + # the runtime owns rolls them back; the worker calls this handler after its + # draft loop, inside the same forward. + self.spec_worker.register_auxiliary_state_handler(self.model.spec_rollback) # ----------------------------------------------------------- multimodal --- # Names the executor reads off a model class by `getattr`; each is silently @@ -3852,10 +5609,11 @@ def __init__(self, model_config: ModelConfig): @property def multimodal_data_device_paths(self) -> list[str]: """Tensors under ``multimodal_data`` the executor moves to CUDA. Absent, it moves - everything, putting `image_grid_thw` on the device where + everything, putting `image_grid_thw` and `video_grid_thw` on the device where `pos_embed_interpolate` and `vision_rope_2d` would each pay a D2H sync.""" return [ "image.pixel_values", + "video.pixel_values_videos", "mrope_config.mrope_position_ids", "mrope_config.mrope_position_deltas", ] @@ -3871,43 +5629,164 @@ def multimodal_token_ids(self) -> torch.Tensor: self._mm_token_ids_cache = torch.tensor(self._mm_token_ids_value, dtype=torch.int32) return self._mm_token_ids_cache - def _multimodal_embeds(self, input_ids, kwargs): - """Run the vision tower for this batch and splice its rows in; returns the - `[T, hidden]` embedding buffer, or None when the step carries no media. + def validate_request(self, request) -> None: + """Admission check, run by the executor on every new request. A media request + the forward would assert on fails here as that request's own error; an assert + inside the step ends the executor under the overlap scheduler. Reads only the + request's host-side payload, and a request without media returns at once. + + Refused: a frame past `_VISION_CALL_MAX_PATCHES` patches (the one bound the + tower keeps), a malformed grid or pixel payload, media whose rows disagree with + the prompt's placeholder count, missing or mis-sized `mrope_position_ids`, and a + request with both images and videos whose `mm_item_order` is missing or does not + list its items.""" + data = request.py_multimodal_data + if not data: + return + grids = {} + for modality, (pixel_key, grid_key) in _MEDIA_KEYS.items(): + media = data.get(modality) + if media is not None: + grids[modality] = self._media_grid(modality, media, pixel_key, grid_key) + if not grids: + return + + mrope_ids = (data.get("mrope_config") or {}).get("mrope_position_ids") + if ( + mrope_ids is None + or mrope_ids.dim() != 3 + or mrope_ids.shape[0] != _QSA_POSITION_AXES + or mrope_ids.shape[2] != request.prompt_len + ): + raise ValueError( + f"a media request needs the [{_QSA_POSITION_AXES}, 1, {request.prompt_len}] " + "mrope_config.mrope_position_ids from the input processor, one position per " + "prompt token; got " + f"{None if mrope_ids is None else tuple(mrope_ids.shape)}" + ) + + if len(grids) > 1: + self._check_item_order(grids, getattr(request, "py_mm_item_order", None)) + + cumsum = data.get("multimodal_embed_mask_cumsum") + if cumsum is not None: + merge_unit = self.model.v_merge * self.model.v_merge + rows = sum(t * h * w for grid in grids.values() for t, h, w in grid) // merge_unit + placeholders = int(cumsum[-1]) + if placeholders != rows: + raise ValueError( + f"the prompt holds {placeholders} image/video placeholder tokens and its " + f"{' and '.join(sorted(grids))} payload produces {rows} embedding rows" + ) + + def _media_grid(self, modality, media, pixel_key, grid_key): + """One modality's `(frames, height, width)` rows, checked on the host.""" + core = self.model + pixel_values, grid_thw = media.get(pixel_key), media.get(grid_key) + if pixel_values is None or grid_thw is None: + raise ValueError( + f"the {modality} payload carries {sorted(media)}; this target reads " + f"{pixel_key!r} and {grid_key!r} from it" + ) + if grid_thw.dim() != 2 or grid_thw.shape[0] == 0 or grid_thw.shape[1] != 3: + raise ValueError( + f"{grid_key} must hold one (t, h, w) row per {modality}; got shape " + f"{tuple(grid_thw.shape)}" + ) + grid = [tuple(int(v) for v in row) for row in grid_thw.tolist()] + for item, (frames, height, width) in enumerate(grid): + if min(frames, height, width) < 1 or height % core.v_merge or width % core.v_merge: + raise ValueError( + f"{modality} {item} has grid {(frames, height, width)}; every entry must " + f"be positive and the sides divisible by spatial_merge_size {core.v_merge}" + ) + if height * width > _VISION_CALL_MAX_PATCHES: + raise ValueError( + f"{modality} {item} has frames of {height}x{width} = {height * width} " + f"patches ({height * core.v_patch}x{width * core.v_patch} pixels); the " + f"vision tower serves frames of at most {_VISION_CALL_MAX_PATCHES} " + "patches, so send it at a lower resolution" + ) + patches = sum(t * h * w for t, h, w in grid) + if tuple(pixel_values.shape) != (patches, core.v_pixel_width): + raise ValueError( + f"{pixel_key} is {tuple(pixel_values.shape)}; its grids describe " + f"{patches} patch rows of {core.v_pixel_width}" + ) + return grid + + @staticmethod + def _check_item_order(grids, order) -> None: + """A request with several modalities lists its items in prompt order.""" + if not order: + raise ValueError( + f"a request carrying {sorted(grids)} needs mm_item_order to splice them in " + "prompt order; the payload holds each modality's rows apart" + ) + listed = {modality: [] for modality in grids} + for entry in order: + try: + listed[entry["modality"]].append(int(entry["index"])) + except (KeyError, TypeError, ValueError) as error: + raise ValueError( + f"mm_item_order entry {entry!r} is not an item of this request, which " + f"carries {sorted(grids)}; entries are {{'modality', 'index'}}" + ) from error + for modality, grid in grids.items(): + if sorted(listed[modality]) != list(range(len(grid))): + raise ValueError( + f"mm_item_order lists {modality} items {listed[modality]}; the payload " + f"carries {len(grid)}" + ) + + def _multimodal_embeds(self, kwargs): + """Run the vision tower for this batch; returns `(rows, mm_token_indices)` (the + tower's `[R, hidden]` rows in packed-batch order and the `[R]` placeholder rows they + replace, which the prefill target splices into its embedding), or None when the step + carries no media. Lives in the shell because the inherited `DecoderModelForCausalLM.forward` drops `multimodal_params` from `**kwargs` on the no-speculation path (a booting, fluent, blind engine). Host-side bookkeeping: `multimodal_params` is a subset of the context requests, so never zip it positionally; - `multimodal_runtime` (chunked prefill / KV reuse only) gives cached and - in-chunk row counts; `mm_token_indices` names every placeholder row.""" + `multimodal_runtime` (chunked prefill / KV reuse only) gives the cached and + in-chunk row counts over the request's images and videos together, in prompt + order; `mm_token_indices` names every placeholder row. + + The tower runs once per modality a request carries, and `_in_prompt_order` + merges the two outputs. Positions are not derived here: every coordinate of + a vision span, a video's temporal one included, is in the payload's + `mrope_config.mrope_position_ids`, which the engine slices into + `position_ids`.""" mm_params = kwargs.get("multimodal_params") or [] if not mm_params: return None - core = self.model rows = [] for params in mm_params: data = getattr(params, "multimodal_data", None) or {} - image = data.get("image") - assert data.get("video") is None, ( - "this target implements the image path only; a video request " - "reached it. The tower itself is per-frame and would run, but " - "the temporal half of the position bookkeeping is not derived " - "here and a wrong one is invisible to every gate" - ) - if image is None: + image, video = data.get("image"), data.get("video") + if image is None and video is None: continue - pixel_values = image.get("pixel_values") - grid_thw = image.get("image_grid_thw") - assert pixel_values is not None and grid_thw is not None, ( - f"the multimodal payload carries {sorted(image)}; this target " - "reads 'pixel_values' and 'image_grid_thw' from it" - ) - assert not grid_thw.is_cuda, ( - "image_grid_thw must stay on the host; see multimodal_data_device_paths" + mrope_ids = (data.get("mrope_config") or {}).get("mrope_position_ids") + assert ( + mrope_ids is not None + and mrope_ids.dim() == 3 + and mrope_ids.shape[0] == _QSA_POSITION_AXES + ), ( + f"a media request needs mrope_config.mrope_position_ids of shape " + f"[{_QSA_POSITION_AXES}, 1, L]; got " + f"{None if mrope_ids is None else tuple(mrope_ids.shape)}" ) - embeds = core.vision_tower(pixel_values, grid_thw) runtime = getattr(params, "multimodal_runtime", None) + if runtime is not None and int(runtime.num_mm_tokens_in_chunk) == 0: + # Every placeholder row of this request lies in another chunk. + continue + towers = {} + if image is not None: + towers["image"] = self._run_tower(image, *_MEDIA_KEYS["image"]) + if video is not None: + towers["video"] = self._run_tower(video, *_MEDIA_KEYS["video"]) + embeds = self._in_prompt_order(params, towers) if runtime is not None: # Chunked prefill or reused prefix: splice only this chunk's slice. start = int(runtime.num_cached_mm_tokens) @@ -3916,12 +5795,10 @@ def _multimodal_embeds(self, input_ids, kwargs): assert int(runtime.total_embeds_in_request) == total, ( f"the executor counted {runtime.total_embeds_in_request} " f"placeholder rows for this request and the tower produced " - f"{total}; the two are derived from the same grid and a " - "disagreement means the merger's row count is wrong" + f"{total} over its {sorted(towers)}; the two are derived from the " + "same grids and a disagreement means the merger's row count is wrong" ) assert 0 <= start <= start + count <= total, (start, count, total) - if count == 0: - continue embeds = embeds[start : start + count] rows.append(embeds) if not rows: @@ -3937,10 +5814,59 @@ def _multimodal_embeds(self, input_ids, kwargs): f"the batch has {mm_token_indices.shape[0]} placeholder rows and the " f"tower produced {vision.shape[0]} embedding rows" ) - h = nn.functional.embedding(input_ids, core.w["embed"]) - # Indexed write: placeholder rows form several disjoint, data-dependent runs. - h.index_copy_(0, mm_token_indices, vision) - return h + return vision, mm_token_indices + + def _run_tower(self, media, pixel_key, grid_key): + """The tower over one modality's payload: `(rows, grid_thw)`, the rows one per + placeholder, item by item in the order the processor emitted them.""" + pixel_values, grid_thw = media.get(pixel_key), media.get(grid_key) + assert pixel_values is not None and grid_thw is not None, ( + f"the multimodal payload carries {sorted(media)}; this target " + f"reads {pixel_key!r} and {grid_key!r} from it" + ) + assert not grid_thw.is_cuda, ( + f"{grid_key} must stay on the host; see multimodal_data_device_paths" + ) + return self.model.vision_tower(pixel_values, grid_thw), grid_thw + + def _in_prompt_order(self, params, towers): + """One request's tower rows in the order its placeholders appear. + + `towers` maps each modality the request carries to `(rows, grid_thw)`; each + grid row is one item of `t*h*w / merge**2` rows. One modality is already in + prompt order. Two are reassembled item by item from `mm_item_order` + (`{"modality", "index"}` per item in prompt order, `index` counting that + modality's items); without it there is no host-side way to tell which + placeholder run is an image and which a video.""" + if len(towers) == 1: + return next(iter(towers.values()))[0] + order = getattr(params, "mm_item_order", None) + assert order, ( + f"a request carrying {sorted(towers)} needs mm_item_order to splice them " + "in prompt order; the payload holds each modality's rows apart" + ) + merge_unit = self.model.v_merge * self.model.v_merge + items = {} + for modality, (embeds, grid_thw) in towers.items(): + counts = [t * h * w // merge_unit for t, h, w in grid_thw.tolist()] + assert sum(counts) == embeds.shape[0], ( + f"the {modality} grid describes {sum(counts)} rows and the tower " + f"produced {embeds.shape[0]}" + ) + items[modality] = torch.split(embeds, counts, 0) + listed = {modality: [] for modality in items} + for entry in order: + assert entry["modality"] in items, ( + f"mm_item_order lists a {entry['modality']!r} item; this request " + f"carries {sorted(items)}" + ) + listed[entry["modality"]].append(int(entry["index"])) + for modality, parts in items.items(): + assert sorted(listed[modality]) == list(range(len(parts))), ( + f"mm_item_order lists {modality} items {listed[modality]}; the " + f"payload carries {len(parts)}" + ) + return torch.cat([items[entry["modality"]][int(entry["index"])] for entry in order], 0) def forward( self, @@ -3960,16 +5886,18 @@ def forward( it would drop it from what the base receives). With one, the shell owns the logits gather (one-model modes set `without_logits`) and the worker needs ungathered hidden states; `position_ids` is passed in the engine's layout - (flattening gives silently wrong draft positions). The vision splice happens - here because the base forward would not pass `multimodal_params` through.""" - spliced = self._multimodal_embeds(input_ids, kwargs) - if spliced is not None: - assert inputs_embeds is None, ( - "the engine supplied inputs_embeds and this step also carries " - "multimodal params; this target owns the splice and cannot " - "reconcile two embedding sources" - ) - inputs_embeds = spliced + (flattening gives silently wrong draft positions). The vision tower runs here + because the base forward would not pass `multimodal_params` through; its rows reach + the core through `begin_step`, and the prefill target splices them.""" + self.model.begin_step( + attn_metadata, + input_ids, + position_ids, + inputs_embeds=inputs_embeds, + lora_params=lora_params, + spec_metadata=spec_metadata, + vision=self._multimodal_embeds(kwargs), + ) if self.spec_worker is None: assert spec_metadata is None, ( "spec_metadata arrived without a spec worker; this engine was " @@ -3989,14 +5917,16 @@ def forward( "the spec worker is built but the engine passed no spec_metadata" ) # `final` is the `[T, hidden]` mix for `lm_head`, `hc` the stream the MTP reads. - final, hc = self.model( - attn_metadata=attn_metadata, - input_ids=input_ids, - position_ids=position_ids, - inputs_embeds=inputs_embeds, - lora_params=lora_params, - spec_metadata=spec_metadata, - **kwargs, + final, hc = self.model.split_spec_output( + self.model( + attn_metadata=attn_metadata, + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + lora_params=lora_params, + spec_metadata=spec_metadata, + **kwargs, + ) ) # `gather_ids`: one row per context request, `runtime_draft_len + 1` per # generation request. @@ -4006,6 +5936,12 @@ def forward( attn_metadata, True, ) + # A prefill CUDA graph pads the step's rows; the worker reshapes generation rows per + # request, so it takes the batch's own. + if attn_metadata.padded_num_tokens is not None: + n = attn_metadata.num_tokens + input_ids, hc = input_ids[:n], hc[:n] + position_ids = position_ids[..., :n] return self.spec_worker( input_ids=input_ids, position_ids=position_ids, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_erf.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_erf.py index 7bfbe1ee84a6..8870a7a5eeea 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_erf.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_erf.py @@ -25,8 +25,9 @@ #: 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) +#: any MoE tile size -- and 8192 the most one tower call carries (32768 patch +#: rows). +MERGER_ROW_COUNTS = (1, 7, 64, 65, 140, 1024, 4096, 8192) _SQRT_2 = math.sqrt(2.0) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_tanh.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_tanh.py index 2f6c5955f1dc..a4ef7a0fe27a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_tanh.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_tanh.py @@ -50,6 +50,14 @@ class GeluTanh(OpWrapper): ) for rows in (1, 2, 7, 64, 1024, 4096) ) + + ( + Cell( + why="block MLP activand at the tower call ceiling, 32768 patch rows x 4304 bf16", + spec=dict( + shape=(32768, INTERMEDIATE_SIZE), dtype=torch.bfloat16, lo=-16.0, hi=16.0 + ), + ), + ) + ( Cell( why="activand built through fc1 at the checkpoint's measured scale", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_canonical_selection.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_canonical_selection.py new file mode 100644 index 000000000000..130518c7111c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_canonical_selection.py @@ -0,0 +1,278 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""QSA block selection, made canonical: one answer per row, whichever tied columns the +selection engine happened to keep. + +Composed from native torch ops: no registered op rewrites a selection, and none of the +top-k engines (`indexer_topk_prefill`, the CuTe DSL kernels, FlashInfer's top-k) takes a +tie-break rule or fixes the order of its output. +""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: This checkpoint's block budget: `indexer_budget` 2048 tokens over +#: `indexer_compress_ratio` 4. +_BLOCK_TOPK = 512 +#: Score-row width the target hands over: a 260-page block table at 32 tokens +#: per block, compressed 4:1. +_SCORE_COLUMNS = 260 * 32 // 4 +#: The most `[rows, cols]` elements one rewrite pass holds at once; a taller matrix is +#: rewritten in row tiles so the fp32 / int32 temporaries stay a few hundred MiB. +_TILE_ELEMENTS = 1 << 24 + + +def _cell(why: str, rows, windows, logits: str, tail: str = "nan", **extra) -> Cell: + return Cell( + why=why, + spec=dict( + rows=rows, + cols=extra.pop("cols", _SCORE_COLUMNS), + index_topk=_BLOCK_TOPK, + windows=windows, + logits=logits, + tail=tail, + selection=extra.pop("selection", "random"), + **extra, + ), + ) + + +def _rewrite_tile(logits, row_starts, row_ends, indices) -> None: + num_rows, num_cols = logits.shape + index_topk = indices.shape[1] + starts = row_starts[:, None] + columns = torch.arange(num_cols, dtype=torch.int32, device=logits.device)[None, :] + # The window's scores, NaN elsewhere: NaN compares false, so nothing outside the + # window is ever counted or kept. + windowed = torch.where( + (columns >= starts) & (columns < row_ends[:, None]), logits, float("nan") + ) + # The k-th value: the smallest score the engine kept. + picked = torch.gather(logits, 1, (indices + starts).clamp_(0, num_cols - 1).long()) + kth = picked.masked_fill_(indices < 0, float("inf")).amin(1, keepdim=True) + at_least = windowed >= kth + tied = windowed == kth + # Keep every column at or above the k-th value except the lowest-index ties past + # k: a tied column is dropped while no more than `surplus` ties lie at or before + # it. A window of at most k has no surplus. + surplus = at_least.sum(1, keepdim=True, dtype=torch.int32) - index_topk + dropped = tied & (torch.cumsum(tied, 1, dtype=torch.int32) <= surplus) + kept_seen = torch.cumsum(at_least ^ dropped, 1, dtype=torch.int32) + # The j-th kept column is the first whose running count reaches j. + slots = torch.arange(1, index_topk + 1, dtype=torch.int32, device=logits.device) + slots = slots.repeat(num_rows, 1) + torch.searchsorted(kept_seen, slots, out_int32=True, out=indices) + indices.sub_(starts) + indices.masked_fill_(slots > kept_seen[:, -1:], -1) + + +class QsaCanonicalSelection(OpWrapper): + """Rewrite an exact top-k selection, in place, into the window's canonical selection. + + For row `r` with window `[s, e) = [row_starts[r], row_ends[r])` of `logits` and + `k = indices.shape[1]`, `indices[r]` holds on entry what a top-k engine wrote: + window-local offsets of an exact top-k of the window, in any order, `-1` past + `min(k, e - s)`. Let `v` be the smallest selected score. On return the row is + + * every window column scoring strictly above `v`, then the highest-index window + columns scoring exactly `v` until there are `k` -- the k largest + `(score, column)` pairs; or the whole window when `e - s <= k`; + * as window-local offsets in ascending order, `-1` in every remaining slot. + + So the row is a function of `(logits, row_starts, row_ends, k)` alone: two engines, + or two launches of one engine, that broke a tie at the k-th value differently produce + the same row. On a row with no tie at the k-th value this is the engine's own set, + sorted. Ties are float equality: `-0.0` and `+0.0` are one tie class. Returns + `indices`. + + Fusion boundary: the rewrite only. Selecting is the engine's call before it; the + engine must be exact -- `v` is read off its picks, and a pick below the true k-th + value would admit more than `k` columns above it. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # The target calls this right after `qsa_topk_select`, at `k = 512` over the + # `[rows, 2080]` score matrix whose columns past each `row_end` are allocator + # garbage (the cells poison them). `selection` names what produced the input row: + # `random` is an exact top-k that keeps a random subset of the tie class at the k-th + # value, in a random order (the adversary); `radix` is the engine the target calls. + CELLS: tuple[Cell, ...] = ( + _cell( + "decode: 1, 8, 128 and 256 rows (MTP3 at batch 32, MTP7 at batch 32)", + (1, 8, 128, 256), + ("positions_random", 2051, 8191), + "relu", + ), + _cell("prefill: 1024 rows from position 4000", 1024, ("chunk", 4000), "relu"), + _cell("prefill: 8192 rows from position 0", 8192, ("chunk", 0), "relu"), + _cell( + "prefill: 8192 rows over an 8192-wide table, rewritten in row tiles", + 8192, + ("chunk", 24000), + "relu_sparse", + cols=8192, + ), + _cell( + "windows 0, 1, 511, 512, 513, 1024, 2047, 2048, 2080", + 18, + ("ends", (0, 1, 511, 512, 513, 1024, 2047, 2048, 2080) * 2), + "relu", + ), + _cell( + "ties: fewer than k positive scores, the zero class straddles k", + 64, + ("ends", (2048,) * 32 + (1500,) * 16 + (2080,) * 16), + "relu_sparse", + ), + _cell( + "ties: two-valued rows", + 32, + ("ends", (2048,) * 16 + (777,) * 16), + "two_valued", + ), + _cell( + "ties: every window column equal", + 16, + ("ends", (2048,) * 8 + (513,) * 8), + "all_equal", + ), + _cell( + "ties: one positive pair tied exactly at the k-th value", + 16, + ("ends", (2048,) * 16), + "pair_at_k", + ), + _cell( + "ties: -0.0 and +0.0 mixed in the zero class", + 16, + ("ends", (2048,) * 16), + "signed_zeros", + ), + _cell( + "tail of +inf / 3e38 / NaN past the window is never kept", + 32, + ("positions_random", 2051, 8191), + "relu_sparse", + tail="poison", + ), + _cell( + "nonzero row_starts: offsets stay window-local", + 8, + ( + "explicit", + (1, 2, 3, 37, 100, 480, 1001, 1500), + (2080, 600, 515, 1061, 612, 2079, 1513, 2080), + ), + "relu_sparse", + tail="poison", + ), + _cell( + "input from the radix engine (qsa_topk_select), zero ties", + 64, + ("positions_random", 2051, 8191), + "relu_sparse", + selection="radix", + ), + _cell( + "CUDA graph: 32 rows captured, replayed with new scores and ends", + 32, + ("positions_random", 0, 8191), + "relu", + graph=True, + ), + ) + + note = """ + Metadata consumed: none. Stateless; every input is an argument, and `indices` is + both the engine's selection read and the canonical row written. + + Preconditions (none checked): + - `indices` holds an EXACT top-k of each window: the smallest selected score is + the window's k-th largest. `qsa_topk_select` is exact; an approximate selection + would be rewritten into a row with more than `k` columns at or above its k-th + pick, truncated to the lowest `k` column indices -- wrong, silently. + - `0 <= row_starts[r] <= row_ends[r] <= logits.shape[1]`. Anything outside the + window is never read for a decision, so garbage, NaN and +inf there are fine + (the cells poison it). + - Scores inside the window are not NaN: a NaN column is never kept, and a NaN pick + empties the row. (The engine returns the identity for a window of at most k + without reading it; this rewrite reads it.) + + Cost: about 30 kernels -- elementwise passes, two int32 prefix sums over + `[rows, cols]` and one sorted search over `[rows, k]` -- with no host read and no + data-dependent shape, so it captures into a CUDA graph and replays with new scores + and windows. It reads the whole `cols`-wide score row, not the live window: a caller + that knows a bound on the windows passes `logits[:, :bound]`, which gives the same + rows. A matrix of more than 2^24 elements is rewritten in row tiles. + """ + + def __call__( + self, + logits: torch.Tensor, + row_starts: torch.Tensor, + row_ends: torch.Tensor, + indices: torch.Tensor, + ) -> torch.Tensor: + num_rows, num_cols = logits.shape + tile = max(1, _TILE_ELEMENTS // max(num_cols, 1)) + for lo in range(0, num_rows, tile): + hi = min(lo + tile, num_rows) + _rewrite_tile(logits[lo:hi], row_starts[lo:hi], row_ends[lo:hi], indices[lo:hi]) + return indices + + def reference( + self, + logits: torch.Tensor, + row_starts: torch.Tensor, + row_ends: torch.Tensor, + indices: torch.Tensor, + ) -> torch.Tensor: + """The k largest `(score, column)` pairs of each window, ascending by column, + `-1`-padded: a top-k over one int64 key per column whose high word orders the + scores as fp32 does and whose low word is the column. `indices` supplies only the + width `k`; its contents are ignored.""" + num_rows, num_cols = logits.shape + index_topk = indices.shape[1] + columns = torch.arange(num_cols, device=logits.device)[None, :] + starts, ends = row_starts.long()[:, None], row_ends.long()[:, None] + in_window = (columns >= starts) & (columns < ends) + # `+ 0.0` folds -0.0 into +0.0, so the two are one tie class. + bits = (logits.float() + 0.0).view(torch.int32).long() + order = torch.where(bits < 0, bits ^ 0x7FFFFFFF, bits) + key = torch.where(in_window, (order << 32) | columns, torch.iinfo(torch.int64).min) + width = min(index_topk, num_cols) + top = torch.topk(key, width, dim=1).indices + taken = torch.gather(in_window, 1, top) + pad = torch.iinfo(torch.int64).max + ordered = torch.sort(torch.where(taken, top, pad), dim=1).values + out = torch.where(ordered == pad, -1, ordered - starts).to(torch.int32) + if width < index_topk: + out = torch.nn.functional.pad(out, (0, index_topk - width), value=-1) + return out + + def is_valid( + self, + logits: torch.Tensor, + row_starts: torch.Tensor, + row_ends: torch.Tensor, + indices: torch.Tensor, + ) -> None: + # A one-element `row_starts` / `row_ends`, or a one-row `indices`, broadcasts over + # every row of `logits` without an error: a [1] `row_starts` against 4 rows + # returns every row framed by row 0's start. + num_rows = logits.shape[0] + assert row_starts.shape == (num_rows,) and row_ends.shape == (num_rows,), ( + f"row_starts and row_ends need one entry per logits row ({num_rows}); got " + f"{tuple(row_starts.shape)} and {tuple(row_ends.shape)}" + ) + assert indices.dim() == 2 and indices.shape[0] == num_rows, ( + f"indices must be [{num_rows}, k]; got {tuple(indices.shape)}" + ) + + +qsa_canonical_selection = QsaCanonicalSelection() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_expand_block_indices.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_expand_block_indices.py index 3bf02f3b3eac..0744a4ca1f2d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_expand_block_indices.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_expand_block_indices.py @@ -17,6 +17,10 @@ _CKPT = dict(compress_ratio=4, token_topk=2048) +#: First positions of 32 MTP verify blocks in the sparse branch (KV past 2048). +_VERIFY_STARTS = tuple(2100 + (499 * g) % 6000 for g in range(32)) + + def _spec(batch, **overrides) -> dict: """A cell spec at this checkpoint's geometry and the engine's metadata dtypes. @@ -25,6 +29,8 @@ def _spec(batch, **overrides) -> dict: ("serving", requests, context_lens) one packed step: context chunks, then two MTP rows per remaining request ("rows", positions) explicit positions, length = position + 1 + ("verify", starts, p) one MTP verify block of p rows per start, + every row's length the block's end """ base = dict( batch=batch, @@ -77,6 +83,24 @@ class QsaExpandBlockIndices(OpWrapper): 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))), + # ---- MTP verify blocks (max_draft_len up to 7): p rows per request, every + # row's length the block's end, so only the last row has length == position + 1. + Cell( + why="verify MTP7: one request's 8 rows, column split", + spec=_spec(("verify", (6000,), 8)), + ), + Cell( + why="verify MTP7: 15 requests x 8 = 120 rows, column split", + spec=_spec(("verify", _VERIFY_STARTS[:15], 8)), + ), + Cell( + why="verify MTP7: 16 requests x 8 = 128 rows, whole-row programs", + spec=_spec(("verify", _VERIFY_STARTS[:16], 8)), + ), + Cell( + why="verify MTP7: 32 requests x 8 = 256 rows, whole-row programs", + spec=_spec(("verify", _VERIFY_STARTS, 8)), + ), # ---- serving steps: sequence_lengths > query_position + 1 inside chunks Cell( why="mixed 1334 rows: 4 chunks + 28 MTP decodes, 32 slots", @@ -160,8 +184,8 @@ class QsaExpandBlockIndices(OpWrapper): 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. + CUDA-graph capturable (captured at 4 and at 256 rows); a replay recomputes + from the live buffers, and the captured row count fixes the launch shape. Hard limits (read from the kernel, not probed): the row offset is `program_id * stride(0)` in int32, so `rows * stride(0) >= 2^31` (4 Mi diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_mtp_reuse_topk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_mtp_reuse_topk.py new file mode 100644 index 000000000000..ad419f0cd502 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_mtp_reuse_topk.py @@ -0,0 +1,362 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""QSA MTP IndexShare: draft-decode selections rebuilt from the row captured at draft-extend.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +# Tier-2 sourcing: a vendored Triton callable reached as ordinary Python. QSA's +# kernel plane has no torch.ops surface. +from tensorrt_llm._torch.attention.backends.sparse.qsa import kernels as _qsa_kernels + +#: This checkpoint's selection width: `indexer_budget` 2048 tokens over +#: `indexer_compress_ratio` 4. +_BLOCK_TOPK = 512 +#: `(slots, this layer's slot)` of the shared buffers: with the MTP layer in the +#: target's cache manager, 12 trunk full-attention layers plus the MTP layer; a +#: separate draft manager owns one slot, the MTP layer's. +_COMBINED_SLOTS = (13, 12) +_DRAFT_SLOTS = (1, 0) + + +def _spec(contexts, **overrides) -> dict: + """A cell spec at this checkpoint's geometry and the engine's metadata dtypes. + + `contexts` says what each request's captured row and current row are: + ("mixed", B) B requests; the first 16 at fixed boundary positions + (empty, short, 511, 512 and full rows), the rest seeded + ("short", B) captured positions <= 2046: fewer than 512 valid groups + ("full", B) captured positions >= 2047: all 512 slots valid + ("explicit", P) one request per captured position in P + ("gaps", pairs) (captured_visible, current_visible) per request, set + directly rather than through positions + A request captured at position `p` (its last accepted row) holds + `(p + 1) // 4` complete groups; its draft row `d` positions later sees + `(p + d + 1) // 4`. `drafted` cycles over the requests. + """ + base = dict( + contexts=contexts, + drafted=(1, 2, 3), + rows_per_request=1, + request_map="identity", + shared_rows=64, + slots=_COMBINED_SLOTS, + block_topk=_BLOCK_TOPK, + layout="engine", + metadata_dtype=torch.int32, + graph_batches=(), + ) + return {**base, **overrides} + + +class QsaMtpReuseTopk(OpWrapper): + """Rebuild each query row's group selection from its request's captured row, + with the groups that completed since the capture written into it. + + Wraps `tensorrt_llm._torch.attention.backends.sparse.qsa.kernels. + triton_qsa_mtp_reuse_topk`: one Triton launch, one program per output row. + For row `r`, with `q = request_indices[r]` and `K = output.shape[1]` + (`block_topk`): + + 1. start from the captured row `shared_indices[q]`: valid group ids as a + prefix, then `-1`; + 2. `captured = shared_visible_blocks[q]`, `current = visible_blocks[r]`, + `fresh = clamp(current - captured, 0, K)`: the groups + `captured .. captured + fresh - 1` completed after the capture; + 3. `filled` = the count of non-negative ids in the captured row; the + fresh ids go, ascending, into slots `start .. start + fresh - 1` with + `start = min(filled, K - fresh)` -- right after the prefix when there + is room, otherwise after the captured row's newest `K - fresh` ids: + the oldest `filled + fresh - K` ids are displaced and the survivors + shift to the front; + 4. every other slot is the captured row's survivors in their order, `-1` + padding included. + + Every slot of `output[:rows]` is written and nothing else; returns `output`. + The inputs are read-only. Nothing is sorted, compacted or de-duplicated and + the ranking is not recomputed: the captured row is taken as ascending (the + target makes it canonical), so the displaced ids are the oldest groups, not the + lowest-scored ones, and the result is ascending again. + + Fusion boundary: the copy-and-append only. Capturing each request's row + into the shared buffers, the per-row request map and visible counts, and + the expansion into token indices plus the incomplete causal tail + (`qsa_expand_block_indices`) are the caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + # ---- single requests: the two behaviours --------------------------- + Cell( + why="1 request, short row: group 100 appended after a 100-id prefix", + spec=_spec(("explicit", (402,)), drafted=(1,), shared_rows=8), + ), + Cell( + why="1 request, full row: group 512 displaces the oldest id, slot 0", + spec=_spec(("explicit", (2050,)), drafted=(1,), shared_rows=8), + ), + Cell( + why="8 requests, one per boundary: 0/1/100/511/512/2047 groups, fresh 0 and 1", + spec=_spec( + ("explicit", (2, 3, 402, 2046, 2050, 8188, 2051, 2047)), + drafted=(1, 1, 1, 1, 1, 3, 3, 1), + shared_rows=8, + ), + ), + # ---- decode batches: one row per request ---------------------------- + Cell(why="decode MTP: 4 requests, mixed rows", spec=_spec(("mixed", 4))), + Cell( + why="decode MTP: 32 requests (envelope), 256-row shared buffer", + spec=_spec(("mixed", 32), shared_rows=256), + ), + Cell( + why="decode MTP: 64 requests, the shared buffer exactly full", + spec=_spec(("mixed", 64), shared_rows=64), + ), + # ---- positions drafted since the capture: draft step i is i past it -- + Cell( + why="32 requests, every one 1 position past its capture", + spec=_spec(("mixed", 32), drafted=(1,)), + ), + Cell( + why="32 requests, every one 2 positions past its capture", + spec=_spec(("mixed", 32), drafted=(2,)), + ), + Cell( + why="32 requests, every one 3 positions past its capture", + spec=_spec(("mixed", 32), drafted=(3,)), + ), + Cell( + why="32 requests, every one 4 positions past its capture", + spec=_spec(("mixed", 32), drafted=(4,)), + ), + Cell( + why="32 requests, every one 5 positions past its capture", + spec=_spec(("mixed", 32), drafted=(5,)), + ), + Cell( + why="32 requests, every one 6 positions past its capture (max_draft_len 7)", + spec=_spec(("mixed", 32), drafted=(6,)), + ), + Cell( + why="8 requests at fresh 2: full, 511-id and short rows gain two groups", + spec=_spec( + ("explicit", (2050, 402, 2046, 8186, 6, 2047, 1022, 2049)), + drafted=(5, 6, 5, 6, 6, 4, 5, 6), + shared_rows=8, + ), + ), + # ---- captured-row fill --------------------------------------------- + Cell( + why="32 short contexts: every captured row below 512 ids", + spec=_spec(("short", 32)), + ), + Cell( + why="32 contexts past 2048 tokens: every captured row full", + spec=_spec(("full", 32)), + ), + Cell( + why="32 short contexts, 4..6 positions past the capture: fresh up to 2", + spec=_spec(("short", 32), drafted=(4, 5, 6)), + ), + Cell( + why="32 full rows, 4..6 positions past the capture: up to 2 oldest ids displaced", + spec=_spec(("full", 32), drafted=(4, 5, 6)), + ), + # ---- addressing: the request map and the slot ----------------------- + Cell( + why="permuted request map, draft manager's single slot", + spec=_spec(("mixed", 32), request_map="permuted", slots=_DRAFT_SLOTS, shared_rows=32), + ), + Cell( + why="3 rows per request (24 rows, 8 requests) share one captured row", + spec=_spec(("mixed", 8), rows_per_request=3, shared_rows=8), + ), + Cell( + why="group gaps 2..600 and current < captured (not reached at cr 4)", + spec=_spec( + ( + "gaps", + ( + (0, 600), + (10, 15), + (500, 1100), + (100, 98), + (2047, 2048), + (300, 812), + (300, 811), + ), + ), + shared_rows=8, + ), + ), + # ---- CUDA graph ----------------------------------------------------- + Cell( + why="CUDA graph: captured at 1, 8 and 32 requests, new-data replay", + spec=_spec(("mixed", 64), graph_batches=(1, 8, 32)), + ), + Cell( + why="CUDA graph, drafted 1..6: captured at 1, 8 and 32 requests, new-data replay", + spec=_spec(("mixed", 32), drafted=(1, 2, 3, 4, 5, 6), graph_batches=(1, 8, 32)), + ), + # ---- dtypes, layouts, geometry -------------------------------------- + Cell( + why="int64 request map and visible counts", + spec=_spec(("mixed", 8), metadata_dtype=torch.int64, shared_rows=8), + ), + Cell( + why="dense standalone tensors instead of engine buffer slices", + spec=_spec(("mixed", 8), layout="dense", shared_rows=8), + ), + Cell( + why="geometry block_topk 8", + spec=_spec(("explicit", (2, 30, 34, 3, 41, 13)), block_topk=8, shared_rows=8), + ), + Cell( + why="geometry block_topk 1024", + spec=_spec(("explicit", (2, 4094, 4098, 402, 8188)), block_topk=1024, shared_rows=8), + ), + Cell(why="zero rows: output returned, no launch", spec=_spec(("mixed", 0))), + ) + + note = """ + Metadata consumed: none; every input is an argument, all int32 at the call + site: + shared_indices `qsa_shared_topk_indices[slot]`, [max_num_sequences, block_topk] + shared_visible_blocks `qsa_shared_topk_visible_blocks[slot]`, [max_num_sequences] + request_indices `qsa_req_idx_per_token[:rows]` + visible_blocks `qsa_visible_blocks[:rows]` + output `qsa_topk_indices[:rows]`, the [max_num_tokens, block_topk] buffer + `slot` is the layer's row in the cache manager's `qsa_shared_topk_slots`. + + It runs in the MTP draft loop's decode steps after the first: no context + rows, one row per request, so rows = requests, and the request map is the + row index. Draft step `i` sits `i` positions past the captured row, so + `max_draft_len` 1 never reaches this op, 3 reaches 1 and 2 positions, and 7 + (the target's ceiling) reaches 1..6. With compress_ratio 4, `fresh = + ((p + 1) % 4 + i) // 4` for a capture at position `p`: 0 or 1 up to i = 3 + and 0, 1 or 2 from i = 4; the gap cell covers wider gaps. In a full row + the fresh ids displace the row's oldest ids (one at fresh 1, two at fresh + 2), and a 511-id row at fresh 2 loses its oldest id. + + Preconditions that are values, so not `is_valid` guards: + - each captured row is prefix-packed (every valid id before every -1), as + `qsa_topk_select` emits it. The append slot comes from the COUNT of + non-negative ids, so a -1 hole puts a fresh id over a valid one and + loses a group silently; + - captured ids are below that row's `shared_visible_blocks`, so a fresh id + never duplicates a captured one; + - every `request_indices[r]` is a row of both shared tensors; + - `shared_visible_blocks[q]` is the visible count of the very row that was + captured, and `visible_blocks[r]` is `(position + 1) // cr` of row `r`. + + Loud rejections: `ValueError` when the shared width differs from the output + width, when `request_indices` or `visible_blocks` does not hold one element + per output row, or on CPU tensors; a 1-D `output` or `shared_indices`; a + width that is not a power of two (Triton compilation). + + Accepted and silently wrong, hence `is_valid`: any non-contiguous operand + (the kernel addresses raw pointers densely and ignores strides); a 3-D + `shared_indices` whose middle extent equals `block_topk` (read as slot 0); + a 2-D `shared_visible_blocks` (the whole `[slots, seqs]` buffer is read as + slot 0's counts). + + Certified: int32 ids, int32 or int64 metadata, `block_topk` 8, 512 and 1024. + No host synchronization and shape-only validation, so it captures into a + CUDA graph; `rows == 0` returns `output` without a launch. + """ + + def __call__( + self, + *, + shared_indices: torch.Tensor, + shared_visible_blocks: torch.Tensor, + request_indices: torch.Tensor, + visible_blocks: torch.Tensor, + output: torch.Tensor, + ) -> torch.Tensor: + return _qsa_kernels.triton_qsa_mtp_reuse_topk( + shared_indices=shared_indices, + shared_visible_blocks=shared_visible_blocks, + request_indices=request_indices, + visible_blocks=visible_blocks, + output=output, + displace_oldest=True, + ) + + def reference( + self, + *, + shared_indices: torch.Tensor, + shared_visible_blocks: torch.Tensor, + request_indices: torch.Tensor, + visible_blocks: torch.Tensor, + output: torch.Tensor, + ) -> torch.Tensor: + """The rebuilt rows as the contract states them, in int64 torch. + + `output` supplies only the shape and dtype; its contents are ignored. + """ + rows, block_topk = output.shape + device = output.device + requests = request_indices.long().reshape(-1) + captured = shared_visible_blocks.long()[requests] + current = visible_blocks.long().reshape(-1) + rebuilt = shared_indices.long()[requests].clone() + + filled = (rebuilt >= 0).sum(dim=1) + fresh = (current - captured).clamp(min=0, max=block_topk) + start = torch.minimum(filled, block_topk - fresh) + slot = torch.arange(block_topk, device=device)[None, :] + # The oldest `filled - start` ids go; the survivors move to the front. + source = slot + (filled - start)[:, None] + rebuilt = torch.where( + source < block_topk, + rebuilt.gather(1, source.clamp(max=block_topk - 1)), + torch.full_like(rebuilt, -1), + ) + rank = slot - start[:, None] + appended = (rank >= 0) & (rank < fresh[:, None]) + rebuilt[appended] = (captured[:, None] + rank)[appended] + return rebuilt.to(output.dtype) + + def is_valid( + self, + *, + shared_indices: torch.Tensor, + shared_visible_blocks: torch.Tensor, + request_indices: torch.Tensor, + visible_blocks: torch.Tensor, + output: torch.Tensor, + ) -> None: + # The kernel takes raw pointers and addresses them densely; each case + # below returns a plausible wrong answer rather than an error. + if shared_indices.dim() != 2 or not shared_indices.is_contiguous(): + raise ValueError( + "shared_indices must be one slot's contiguous [num_seqs, block_topk] " + f"rows; got shape {tuple(shared_indices.shape)}, " + f"stride {shared_indices.stride()}" + ) + if shared_visible_blocks.dim() != 1 or not shared_visible_blocks.is_contiguous(): + raise ValueError( + "shared_visible_blocks must be one slot's contiguous [num_seqs] counts; " + f"got shape {tuple(shared_visible_blocks.shape)}, " + f"stride {shared_visible_blocks.stride()}" + ) + for name, tensor in ( + ("request_indices", request_indices), + ("visible_blocks", visible_blocks), + ("output", output), + ): + if not tensor.is_contiguous(): + raise ValueError( + f"{name} must be contiguous; got shape {tuple(tensor.shape)}, " + f"stride {tensor.stride()}" + ) + + +qsa_mtp_reuse_topk = QsaMtpReuseTopk() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_scores.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_scores.py index d0dbbfa1c66a..5c626cf40ae7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_scores.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_scores.py @@ -71,6 +71,12 @@ def _small(rows: int, requests: int, **overrides) -> dict: ) +#: The same plane under an NVFP4 KV pool: the QSA position page joins each physical +#: page after the index buffers (`page_pad` elements), so the trunk's view is +#: `[pages, 32, 1, 128]` at page stride 12 x 4096 + 192 = 49344. +_TRUNK_POOL_NVFP4 = dict(pool_buffers=12, page_pad=192) + + class QsaPagedIndexScores(OpWrapper): """Score every causally visible compressed block of each query row. @@ -106,6 +112,34 @@ class QsaPagedIndexScores(OpWrapper): ARCHS = frozenset({Arch.SM_103}) CELLS: tuple[Cell, ...] = ( + # ---- the NVFP4-pool plane: page stride 49344, a kernel of its own ---- + Cell( + why="NVFP4 trunk plane, page stride 49344: decode MTP 4 requests x 2 rows", + spec=_spec( + layout="decode", + gen_positions=(7000, 5000, 4200, 2100), + rows_per_request=2, + slots=32, + pages=_WIDTH, + only_visible_blocks=True, + context_rows=False, + branch="row_wise", + **_TRUNK_POOL_NVFP4, + ), + ), + Cell( + why="NVFP4 trunk plane, page stride 49344: prefill 300 rows from 0", + spec=_spec( + layout="capture", + lengths=(300,), + slots=8, + pages=_WIDTH, + only_visible_blocks=True, + context_rows=True, + branch="query_tile", + **_TRUNK_POOL_NVFP4, + ), + ), # ---- the target's call shapes: [*, 260] table, 2080 columns ------- Cell( why="decode: 1 row deep in history, [32, 260] table, visible-only", @@ -159,6 +193,45 @@ class QsaPagedIndexScores(OpWrapper): branch="row_wise", ), ), + Cell( + why="decode MTP7: 1 request x 8 rows across a page edge, visible-only", + spec=_spec( + layout="decode", + gen_positions=(6012,), + rows_per_request=8, + slots=32, + pages=_WIDTH, + only_visible_blocks=True, + context_rows=False, + branch="row_wise", + ), + ), + Cell( + why="decode MTP7: 32 slots x 8 rows = 256 rows, visible-only", + spec=_spec( + layout="decode", + gen_positions=tuple(100 + (499 * g) % 8000 for g in range(32)), + rows_per_request=8, + slots=32, + pages=_WIDTH, + only_visible_blocks=True, + context_rows=False, + branch="row_wise", + ), + ), + Cell( + why="decode MTP7: 32 slots x 8 rows = 256 rows, -inf padded", + spec=_spec( + layout="decode", + gen_positions=tuple(100 + (499 * g) % 8000 for g in range(32)), + rows_per_request=8, + slots=32, + pages=_WIDTH, + only_visible_blocks=False, + context_rows=False, + branch="row_wise", + ), + ), Cell( why="capture: 8191 rows, one request from 0, visible-only", spec=_spec( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_store.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_store.py index 341287f03b95..bef6cf362071 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_store.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_store.py @@ -34,6 +34,10 @@ _DECODE = ((0, 1), (1, 1), (31, 1), (32, 1), (33, 1), (63, 1), (64, 1), (95, 1)) +#: 32 requests of one MTP7 verify block each (8 rows), at starts that put some blocks across a +#: page edge and some inside one. +_MTP7_32 = tuple((29 + 61 * i, 8) for i in range(32)) + def _spec(segments, pages, **overrides) -> dict: """`segments` is `(first position, rows)` per request, in row order.""" @@ -81,6 +85,14 @@ class QsaPagedIndexStore(OpWrapper): 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="decode MTP7: 1 request x 8 rows, a block straddling a page", + spec=_spec(((28, 8),), pages=8), + ), + Cell( + why="decode MTP7: 32 requests x 8 rows (256), blocks straddling pages", + spec=_spec(_MTP7_32, pages=1200, slots=40), + ), Cell( why="prefill: 1024 rows, 4 requests x 256 walking 8 pages each", spec=_spec(((0, 256),) * 4, pages=32), @@ -109,6 +121,18 @@ class QsaPagedIndexStore(OpWrapper): 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="NVFP4-KV side planes: index and position planes as strided views of one " + "per-slot buffer holding 12 index-K pages and the position page", + spec=_spec( + ((0, 1), (31, 1), (32, 1), (33, 40)), + pages=8, + cache_layout="shared_page", + layers=12, + page_stride=49344, + position_strides=(24672, 3, 1), + ), + ), 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"), @@ -151,13 +175,21 @@ class QsaPagedIndexStore(OpWrapper): - duplicate destinations both store, unordered; - coordinates wider than the plane's dtype wrap (int64 `2^31 + 5` stores as -2147483643); floats truncate toward zero. + Under an NVFP4 KV pool the QSA cache manager packs each page's index-K + roles and then its position page into one per-slot buffer, so both planes + are strided views of one allocation: the trunk's index plane is + `[pages, 32, 1, 128]` at page stride 49344 (12 x 4096 + 192) with the + position plane `[pages, 32, 3]` at (24672, 3, 1); the cell checks that the + other index-K roles of each page survive the call. + 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. + rows. CUDA-graph capturable (captured at 4 decode rows and at 256, 32 MTP7 + verify blocks); 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. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store.py index 06796397a5f4..8caa29166a2b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store.py @@ -17,7 +17,9 @@ #: 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. +#: max_blocks_per_seq). `pages` is the pool, independent of the batch. `values` +#: fills the rows: `every_bf16` is all 65536 bf16 bit patterns (needs 128 rows +#: of 2 x 256). _CKPT = dict( hkv=2, head_dim=256, @@ -32,10 +34,17 @@ k_layout="qkv_slice", table_layout="dense", hole=None, + values="randn", ) +_F8 = torch.float8_e4m3fn + _DECODE = ((0, 1), (31, 1), (32, 1), (33, 1), (100, 1), (255, 1), (7, 1), (64, 1)) +#: 32 requests of one MTP7 verify block each (8 rows), at starts that put some blocks across a +#: page edge and some inside one. +_MTP7_32 = tuple((29 + 61 * i, 8) for i in range(32)) + def _spec(segments, pages, **overrides) -> dict: """`segments` is `(first position, rows)` per request, in row order.""" @@ -61,10 +70,11 @@ class QsaPagedKvStore(OpWrapper): `logical_positions` is absolute, so prefill, decode and MTP rows are the same call. The store is a move, converted once to the pool's dtype (bitwise `Tensor.to`; an fp8-e4m3 pool saturates at 448 and takes no - scale). A row that resolves to no slot -- negative request, request past - the table, negative position, page column past the table, or a table - entry that is negative (`-1`, unallocated) or past the pool -- is skipped - without error. Returns `(k_cache, v_cache)`, both written in place. + scale, so a reader treats it as kv scale 1.0). A row that resolves to no + slot -- negative request, request past the table, negative position, page + column past the table, or a table entry that is negative (`-1`, + unallocated) or past the pool -- is skipped without error. Returns + `(k_cache, v_cache)`, both written in place. Fusion boundary: the scatter only. The K/V projection, norm and RoPE of K, the QSA index plane, selection and attention are the caller's; the @@ -92,7 +102,37 @@ class QsaPagedKvStore(OpWrapper): ), Cell( why="fp8-e4m3 pool: the row rounded once", - spec=_spec(_DECODE, pages=64, cache_dtype=torch.float8_e4m3fn), + spec=_spec(_DECODE, pages=64, cache_dtype=_F8), + ), + # The prefill target's call over an fp8 pool (kv_cache_dtype=fp8, unit scale). + Cell( + why="fp8 prefill: 1024 rows, 4 requests x 256 walking 8 pages each", + spec=_spec(((0, 256),) * 4, pages=160, cache_dtype=_F8), + ), + Cell( + why="fp8 mixed: two context chunks and two decode rows, 66 rows", + spec=_spec(((0, 40), (0, 24), (200, 1), (7, 1)), pages=64, cache_dtype=_F8), + ), + Cell( + why="fp8 decode MTP: 4 requests x 2 rows, a pair straddling a page", + spec=_spec(((0, 2), (31, 2), (63, 2), (95, 2)), pages=32, cache_dtype=_F8), + ), + # MTP7 verify blocks: 8 rows per request, so one block spans up to two pages. + Cell( + why="decode MTP7: 1 request x 8 rows, a block straddling a page", + spec=_spec(((28, 8),), pages=8), + ), + Cell( + why="decode MTP7: 32 requests x 8 rows (256), blocks straddling pages", + spec=_spec(_MTP7_32, pages=1200, slots=40), + ), + Cell( + why="fp8 decode MTP7: 32 requests x 8 rows (256), blocks straddling pages", + spec=_spec(_MTP7_32, pages=1200, slots=40, cache_dtype=_F8), + ), + Cell( + why="fp8 conversion: all 65536 bf16 bit patterns, as torch casts them", + spec=_spec(((0, 128),), pages=8, cache_dtype=_F8, values="every_bf16"), ), Cell( why="fp16 rows into a bf16 pool: one conversion", @@ -182,8 +222,9 @@ class QsaPagedKvStore(OpWrapper): 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. + layout or pool size compiles its own kernel. CUDA-graph capturable + (captured at 8 decode rows and at 256, 32 MTP7 verify blocks); a replay + follows positions rewritten in place. Hard limits: both metadata values are cast to int32 inside the kernel (`.to(tl.int32)`); an int64 position at `2^32 + 5` was probed and lands diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store_nvfp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store_nvfp4.py new file mode 100644 index 000000000000..06b54e6a84d0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store_nvfp4.py @@ -0,0 +1,620 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""QSA sparse path's paged K/V append into an NVFP4 pool: quantize and scatter in one launch.""" + +from __future__ import annotations + +from typing import Tuple + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Values per e4m3 block scale, and e2m1 values per stored byte. +_SF_VEC = 16 +_PER_BYTE = 2 +#: Tokens of a V scale group: V's scales are stored `[tpb / 4, D / 16, 4]` per head page. +_V_SF_QUAD = 4 +#: e2m1's largest magnitude: a block's scale maps its amax onto it. +_E2M1_MAX = 6.0 +#: What `cvt.rn.satfinite.e2m1x2.f32` makes of a NaN of either sign: +6. A NaN reaches +#: it as 0 x inf, from a zero in a block whose nonzero amax rounds its scale to 0. +_E2M1_NAN_CODE = 0x7 +#: The round-to-nearest-even decision points between neighbouring e2m1 magnitudes: +#: `(threshold, inclusive)` steps the code up when `|v| > t`, or `|v| >= t` where the +#: upper neighbour is the even code. +_E2M1_STEPS = ( + (0.25, False), + (0.75, True), + (1.25, False), + (1.75, True), + (2.5, False), + (3.5, True), + (5.0, False), +) + +#: This checkpoint's K/V geometry and the engine's QSA table. `heads_q` places k and v as +#: head slices of the packed `[T, 24 + 2 + 2, 256]` QKV row; `global_k` / `global_v` are +#: `kv_scale_orig_quant[1]` / `[2]`; `slot_pages` role pages of one lifecycle slot hold +#: the call's planes, layer `layer`'s K and V among them (the trunk pool holds 24). +_CKPT = dict( + hkv=2, + head_dim=256, + heads_q=24, + tokens_per_block=32, + slots=32, + width=260, + row_dtype=torch.bfloat16, + request_dtype=torch.int32, + position_dtype=torch.int64, + global_k=1.0, + global_v=1.0, + values="randn", + hole=None, + graph_batches=(), + thop_append=False, + slot_pages=2, + layer=0, +) + +_DECODE = ((0, 1), (31, 1), (32, 1), (33, 1), (100, 1), (255, 1), (7, 1), (64, 1)) + + +def _spec(segments, pages, **overrides) -> dict: + """`segments` is `(first position, rows)` per request, in row order.""" + return {**_CKPT, "segments": segments, "pages": pages, **overrides} + + +@triton.jit +def _rcp_approx(x): + """`rcp.approx.ftz.f32`, the reciprocal the engine's FP4 quantizer uses.""" + return tl.inline_asm_elementwise( + "rcp.approx.ftz.f32 $0, $1;", "=f,f", [x], dtype=tl.float32, is_pure=True, pack=1 + ) + + +@triton.jit +def _e2m1x2(lo, hi): + """One byte of two e2m1 codes, `lo` in the low nibble (`cvt.rn.satfinite`).""" + return tl.inline_asm_elementwise( + "{ .reg .b8 t; .reg .b16 s; cvt.rn.satfinite.e2m1x2.f32 t, $2, $1; " + "mov.b16 s, {t, t}; cvt.u32.u16 $0, s; }", + "=r,f,f", + [lo, hi], + dtype=tl.int32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def _quantize_row( + src, + stride_dim, + global_scale, + HEAD_DIM: tl.constexpr, +): + """One head row: (packed e2m1 bytes [D/2], e4m3 scale bytes [D/16]).""" + pairs = tl.arange(0, HEAD_DIM // 2) + lo = tl.load(src + (2 * pairs) * stride_dim).to(tl.float32) + hi = tl.load(src + (2 * pairs + 1) * stride_dim).to(tl.float32) + pair_max = tl.maximum(tl.abs(lo), tl.abs(hi)) + block_max = tl.max(tl.reshape(pair_max, [HEAD_DIM // 16, 8]), axis=1) + sf_value = global_scale * (block_max * _rcp_approx(tl.full([HEAD_DIM // 16], 6.0, tl.float32))) + sf_code = sf_value.to(tl.float8e4nv) + sf_rounded = sf_code.to(tl.float32) + out_scale = tl.where( + block_max != 0.0, + _rcp_approx(sf_rounded * _rcp_approx(tl.zeros_like(sf_rounded) + global_scale)), + 0.0, + ) + lo_scaled = tl.reshape( + tl.reshape(lo, [HEAD_DIM // 16, 8]) * out_scale[:, None], [HEAD_DIM // 2] + ) + hi_scaled = tl.reshape( + tl.reshape(hi, [HEAD_DIM // 16, 8]) * out_scale[:, None], [HEAD_DIM // 2] + ) + packed = (_e2m1x2(lo_scaled, hi_scaled) & 0xFF).to(tl.uint8) + return packed, sf_code.to(tl.uint8, bitcast=True) + + +@triton.jit +def _qsa_paged_kv_store_nvfp4_kernel( + k, + v, + k_cache, + v_cache, + k_scale_cache, + v_scale_cache, + request_indices, + logical_positions, + block_table, + kv_scale_orig_quant, + num_pages, + num_requests, + table_width, + k_stride_row, + k_stride_head, + k_stride_dim, + v_stride_row, + v_stride_head, + v_stride_dim, + kc_stride_page, + kc_stride_head, + kc_stride_token, + vc_stride_page, + vc_stride_head, + vc_stride_token, + ks_stride_page, + ks_stride_head, + vs_stride_page, + vs_stride_head, + bt_stride_request, + bt_stride_page, + HEAD_DIM: tl.constexpr, + TOKENS_PER_BLOCK: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + head = tl.program_id(1).to(tl.int64) + request = tl.load(request_indices + row).to(tl.int64) + logical = tl.load(logical_positions + row).to(tl.int64) + column = logical // TOKENS_PER_BLOCK + token = logical % TOKENS_PER_BLOCK + lookup = (request >= 0) & (request < num_requests) & (logical >= 0) & (column < table_width) + page = tl.load( + block_table + request * bt_stride_request + column * bt_stride_page, mask=lookup, other=-1 + ).to(tl.int64) + valid = (page >= 0) & (page < num_pages) + page = tl.maximum(page, 0) + + k_data, k_sf = _quantize_row( + k + row * k_stride_row + head * k_stride_head, + k_stride_dim, + tl.load(kv_scale_orig_quant + 1), + HEAD_DIM, + ) + v_data, v_sf = _quantize_row( + v + row * v_stride_row + head * v_stride_head, + v_stride_dim, + tl.load(kv_scale_orig_quant + 2), + HEAD_DIM, + ) + pairs = tl.arange(0, HEAD_DIM // 2) + blocks = tl.arange(0, HEAD_DIM // 16) + tl.store( + k_cache + page * kc_stride_page + head * kc_stride_head + token * kc_stride_token + pairs, + k_data, + mask=valid, + ) + tl.store( + v_cache + page * vc_stride_page + head * vc_stride_head + token * vc_stride_token + pairs, + v_data, + mask=valid, + ) + # K scales token-major within the head page; V scales in four-token groups. + tl.store( + k_scale_cache + + page * ks_stride_page + + head * ks_stride_head + + token * (HEAD_DIM // 16) + + blocks, + k_sf, + mask=valid, + ) + tl.store( + v_scale_cache + + page * vs_stride_page + + head * vs_stride_head + + (token // 4) * (4 * (HEAD_DIM // 16)) + + blocks * 4 + + token % 4, + v_sf, + mask=valid, + ) + + +def _store_nvfp4( + *, + k: torch.Tensor, + v: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + k_scale_cache: torch.Tensor, + v_scale_cache: torch.Tensor, + request_indices: torch.Tensor, + logical_positions: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + kv_scale_orig_quant: torch.Tensor, +) -> None: + """Quantize rows of K and V to NVFP4 and store them into their pages (see the entry).""" + if k.ndim != 3 or v.shape != k.shape: + raise ValueError("NVFP4 K/V store: k and v must be matching [rows, heads, dim]") + rows, heads, head_dim = k.shape + if k.dtype not in (torch.bfloat16, torch.float16) or v.dtype != k.dtype: + raise ValueError("NVFP4 K/V store: k and v must both be BF16 or both FP16") + if head_dim < 32 or head_dim & (head_dim - 1): + raise ValueError("NVFP4 K/V store: head_dim must be a power of two of at least 32") + if tokens_per_block <= 0 or tokens_per_block % _V_SF_QUAD: + raise ValueError("NVFP4 K/V store: tokens_per_block must be a positive multiple of 4") + data_shape = (k_cache.shape[0], heads, tokens_per_block, head_dim // _PER_BYTE) + scale_shape = (k_cache.shape[0], heads, tokens_per_block, head_dim // _SF_VEC) + for name, cache, shape in ( + ("k_cache", k_cache, data_shape), + ("v_cache", v_cache, data_shape), + ("k_scale_cache", k_scale_cache, scale_shape), + ("v_scale_cache", v_scale_cache, scale_shape), + ): + if cache.dtype != torch.uint8 or tuple(cache.shape) != shape: + raise ValueError( + f"NVFP4 K/V store: {name} must be uint8 {shape}, got {cache.dtype} " + f"{tuple(cache.shape)}" + ) + if cache.stride(3) != 1: + raise ValueError(f"NVFP4 K/V store: {name} must be contiguous along its last dim") + for name, cache in (("k_scale_cache", k_scale_cache), ("v_scale_cache", v_scale_cache)): + if cache.stride(2) != head_dim // _SF_VEC: + raise ValueError(f"NVFP4 K/V store: each head page of {name} must be contiguous") + if block_table.ndim != 2 or block_table.dtype != torch.int32: + raise ValueError("NVFP4 K/V store: block_table must be a 2-D int32 tensor") + for name, meta in ( + ("request_indices", request_indices), + ("logical_positions", logical_positions), + ): + if meta.shape != (rows,) or meta.dtype not in (torch.int32, torch.int64): + raise ValueError(f"NVFP4 K/V store: {name} must be one int32/int64 entry per row") + if not meta.is_contiguous(): + raise ValueError(f"NVFP4 K/V store: {name} must be contiguous") + if ( + kv_scale_orig_quant.dtype != torch.float32 + or kv_scale_orig_quant.shape != (3,) + or not kv_scale_orig_quant.is_contiguous() + ): + raise ValueError("NVFP4 K/V store: kv_scale_orig_quant must be a contiguous fp32 [3]") + operands = (k, v, k_cache, v_cache, k_scale_cache, v_scale_cache, request_indices) + operands += (logical_positions, block_table, kv_scale_orig_quant) + if not all(t.is_cuda and t.device == k.device for t in operands): + raise ValueError("NVFP4 K/V store: every operand must be on k's CUDA device") + if rows == 0: + return + _qsa_paged_kv_store_nvfp4_kernel[(rows, heads)]( + k, + v, + k_cache, + v_cache, + k_scale_cache, + v_scale_cache, + request_indices, + logical_positions, + block_table, + kv_scale_orig_quant, + k_cache.shape[0], + block_table.shape[0], + block_table.shape[1], + k.stride(0), + k.stride(1), + k.stride(2), + v.stride(0), + v.stride(1), + v.stride(2), + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(2), + v_cache.stride(0), + v_cache.stride(1), + v_cache.stride(2), + k_scale_cache.stride(0), + k_scale_cache.stride(1), + v_scale_cache.stride(0), + v_scale_cache.stride(1), + block_table.stride(0), + block_table.stride(1), + HEAD_DIM=head_dim, + TOKENS_PER_BLOCK=tokens_per_block, + num_warps=4, + ) + + +def _e2m1_codes(values: torch.Tensor) -> torch.Tensor: + """fp32 -> e2m1 codes 0..15, round to nearest even, saturating at 6; a NaN + (0 x inf, in a block whose scale underflows to zero) converts to +6.""" + magnitude = values.abs() + code = torch.zeros(values.shape, dtype=torch.int32, device=values.device) + for threshold, inclusive in _E2M1_STEPS: + code += (magnitude >= threshold if inclusive else magnitude > threshold).to(torch.int32) + code |= torch.signbit(values).to(torch.int32) << 3 + return torch.where(torch.isnan(values), torch.full_like(code, _E2M1_NAN_CODE), code) + + +def _quantize_reference(rows: torch.Tensor, global_scale: torch.Tensor): + """[n, D] rows -> (packed bytes [n, D/2], e4m3 scale bytes [n, D/16]), in plain torch.""" + n, head_dim = rows.shape + x = rows.float().reshape(n, head_dim // _SF_VEC, _SF_VEC) + block_max = x.abs().amax(dim=-1) + sixth = torch.tensor(_E2M1_MAX, dtype=torch.float32, device=rows.device).reciprocal() + scale_code = (global_scale * (block_max * sixth)).to(torch.float8_e4m3fn) + rounded = scale_code.float() + out_scale = torch.where( + block_max != 0, + (rounded * global_scale.reciprocal()).reciprocal(), + torch.zeros_like(rounded), + ) + codes = _e2m1_codes(x * out_scale[..., None]).reshape(n, head_dim // _PER_BYTE, _PER_BYTE) + packed = (codes[..., 0] | (codes[..., 1] << 4)).to(torch.uint8) + return packed, scale_code.view(torch.uint8) + + +class QsaPagedKvStoreNvfp4(OpWrapper): + """Quantize every row's K and V to NVFP4 and store them into the request's page. + + A Triton kernel defined in this file, one launch over `(rows, Hkv)`. For row t + and KV head h, with `page = block_table[request_indices[t], logical_positions[t] + // tpb]` and `s = logical_positions[t] % tpb`, each 16-value block `b` of + `k[t, h]` (and alike for v, with `kv_scale_orig_quant[2]`) is stored as:: + + g = kv_scale_orig_quant[1] + sf = e4m3(g * (amax_b * rcp(6))) # one byte + code_i = e2m1(k[t, h, i] * rcp(float(sf) * rcp(g))) # 0 when amax_b == 0 + k_cache[page, h, s, i // 2] nibble i % 2 (low nibble = even i) + k_scale_cache[page, h, s, b] = sf # token-major + v_scale_cache[page, h, (s // 4) * 4 * (D / 16) + 4 * b + s % 4] # four-token groups + + `rcp` is `rcp.approx.ftz.f32`, `e4m3` and `e2m1` round to nearest even and + saturate. This is the attention op's own NVFP4 append (`cvt_warp_fp16_to_fp4` + behind `quantizeAndWriteFP4KVCache`), the same bits for the same row: a + token's stored value does not depend on which path wrote it. A row that + resolves to no slot -- negative request, request past the table, negative + position, page column past the table, a table entry that is negative or past + the pool -- is skipped. Returns the four caches, written in place. + + Fusion boundary: quantization and the scatter. The projection, norm and RoPE + of K are the caller's; the QSA index plane is not touched. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode: 8 rows at aligned, mid-page and page-crossing positions", + spec=_spec(_DECODE, pages=64), + ), + Cell( + why="decode MTP draft 1: 4 requests x 2 rows, a pair straddling a page", + spec=_spec(((0, 2), (31, 2), (63, 2), (95, 2)), pages=32), + ), + Cell( + why="decode MTP draft 7: 4 requests x 8 rows across page edges", + spec=_spec(((0, 8), (28, 8), (60, 8), (1000, 8)), pages=64), + ), + 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="every finite bf16 value, K in bit-pattern order and V shuffled across blocks", + spec=_spec(((0, 128),), pages=8, values="finite_bf16"), + ), + Cell( + why="zeros, -0.0 and a 900.0 outlier inside blocks: zero scales and saturation", + spec=_spec(_DECODE, pages=64, values="edge"), + ), + Cell( + why="second-level scales g_k 0.37, g_v 2.5 (unit scale is what the target binds)", + spec=_spec(((0, 40), (0, 24), (200, 1), (7, 1)), pages=64, global_k=0.37, global_v=2.5), + ), + Cell( + why="the attention op's own append over the same rows (context, chunk, decode " + "P=1 and P=3) at unit second-level scales: data and both scale layouts bitwise equal", + spec=_spec(((0, 115), (0, 101), (0, 36)), pages=96, thop_append=True), + ), + Cell( + why="the attention op's own append at second-level scales 0.37 and 2.5", + spec=_spec( + ((0, 115), (0, 101), (0, 36)), + pages=96, + thop_append=True, + global_k=0.37, + global_v=2.5, + ), + ), + Cell( + why="the engine's trunk layout: layer 11's planes among the 24 role pages of a slot", + spec=_spec(((0, 64), (1000, 2), (31, 2)), pages=96, slot_pages=24, layer=11), + ), + # Triton compiles once per class of each plain-int argument -- equal to 1, a multiple + # of 16, or anything else. The pool's slot count and the table's height are the two a + # served call varies; the cells above cover (multiple of 16, multiple of 16). + Cell( + why="specialization: slot count other, table height 32", + spec=_spec(_DECODE, pages=69), + ), + Cell( + why="specialization: slot count other, table height 1", + spec=_spec(((0, 1),), pages=69, slots=1), + ), + Cell( + why="specialization: slot count other, table height 8", + spec=_spec(_DECODE, pages=69, slots=8), + ), + Cell( + why="specialization: slot count a multiple of 16, table height 1", + spec=_spec(((0, 1),), pages=64, slots=1), + ), + Cell( + why="specialization: slot count a multiple of 16, table height 8", + spec=_spec(_DECODE, pages=64, slots=8), + ), + Cell( + why="decode MTP graph: G = 1, 3 and 8 requests x 8 rows, replayed bitwise equal to eager", + spec=_spec( + tuple((29 + 61 * i, 8) for i in range(8)), + pages=256, + graph_batches=(1, 3, 8), + ), + ), + 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)), + 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 the + step's `qsa_req_idx_per_token` (int32) and `qsa_logical_positions` (int64), + `qsa_block_table` unsliced (raw pool page ids, `-1` unallocated), the layer's + K and V planes of the QSA manager's `get_buffers(i, "HND")` (uint8 `[slots, + Hkv, tpb, D/2]`), the two planes of `get_block_scale_buffers(i, "HND")` + (uint8 `[slots, Hkv, tpb, D/16]`), and its own `kv_scale_orig_quant` -- the + same fp32 `[3]` tensor its dense attention call binds, so both writers apply + one second-level scale. + + Preconditions: + - destination slots are unique across the batch: two rows naming one + (page, slot) both store, unordered. + - the dense path's reader assumes this layout: e2m1 nibbles HND per page, + K scales linear, V scales in four-token groups. The attention op's + append, run over the same rows, leaves bitwise the same pool (cells), at + unit and non-unit second-level scales. + The callable rejects loudly (ValueError): non-matching or non-3-D k/v, a row + dtype outside bf16/fp16, a head_dim that is not a power of two >= 32, a + tokens_per_block that is not a multiple of 4, caches that are not uint8 of + the stated shapes or whose last dim or scale head page is strided, a non-2-D + or non-int32 table, metadata that is not one contiguous int32/int64 entry per + row, a `kv_scale_orig_quant` that is not a contiguous fp32 `[3]`, operands on + another device. `rows == 0` returns before launching. + + Every offset is formed in 64 bits: rows, heads, requests, positions and pages + are widened before any product. CUDA-graph capturable; a replay follows rows, + pages and positions rewritten in place. + + Compiled variants: the page count (`num_pages`, the pool's slot count), the + table's height (`num_requests`) and width and every stride are plain-int + arguments, and Triton compiles the kernel once per class of each -- equal to + 1, a multiple of 16, anything else -- besides once per head_dim and + tokens_per_block. A served call puts the width and each stride in one fixed + class (the k/v head slices of the packed QKV row, the planes' slot, head and + token strides, the contiguous table's) and varies the page count and the + table height; the cells run both reachable page-count classes against every + table-height class. + + Not finite values: a NaN or +-inf in a block is not certified (the e2m1 and + e4m3 conversions saturate, the block maximum is the kernel's `max`). + """ + + def __call__( + self, + k: torch.Tensor, + v: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + k_scale_cache: torch.Tensor, + v_scale_cache: torch.Tensor, + request_indices: torch.Tensor, + logical_positions: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + kv_scale_orig_quant: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + _store_nvfp4( + k=k, + v=v, + k_cache=k_cache, + v_cache=v_cache, + k_scale_cache=k_scale_cache, + v_scale_cache=v_scale_cache, + request_indices=request_indices, + logical_positions=logical_positions, + block_table=block_table, + tokens_per_block=tokens_per_block, + kv_scale_orig_quant=kv_scale_orig_quant, + ) + return k_cache, v_cache, k_scale_cache, v_scale_cache + + def reference( + self, + k: torch.Tensor, + v: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + k_scale_cache: torch.Tensor, + v_scale_cache: torch.Tensor, + request_indices: torch.Tensor, + logical_positions: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + kv_scale_orig_quant: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """The four caches after the store, as fresh tensors; the inputs are untouched. + + Quantization in fp32 torch with IEEE reciprocals where the kernel uses + `rcp.approx.ftz`; the two agree on every operand the cells drive (6.0, the + second-level scales, and every e4m3 scale value), which the cells' bitwise + gate is what pins. + """ + out = [ + t.clone(memory_format=torch.contiguous_format) + for t in (k_cache, v_cache, k_scale_cache, v_scale_cache) + ] + rows, heads, head_dim = k.shape + requests = request_indices.long() + positions = logical_positions.long() + columns = positions.div(tokens_per_block, rounding_mode="floor") + table_rows, table_width = block_table.shape + ok = (requests >= 0) & (requests < table_rows) & (positions >= 0) & (columns < table_width) + pages = block_table[ + requests.clamp(0, table_rows - 1), columns.clamp(0, table_width - 1) + ].long() + ok &= (pages >= 0) & (pages < k_cache.shape[0]) + page, slot = pages[ok], (positions % tokens_per_block)[ok] + n = int(page.numel()) + if n == 0: + return tuple(out) + groups = head_dim // _SF_VEC + quad_base = (slot // _V_SF_QUAD) * (_V_SF_QUAD * groups) + slot % _V_SF_QUAD + for src, data, scales, global_scale, linear in ( + (k, out[0], out[2], kv_scale_orig_quant[1], True), + (v, out[1], out[3], kv_scale_orig_quant[2], False), + ): + packed, sf = _quantize_reference(src[ok].reshape(n * heads, head_dim), global_scale) + data[page, :, slot, :] = packed.reshape(n, heads, head_dim // _PER_BYTE) + sf = sf.reshape(n, heads, groups) + if linear: + scales[page, :, slot, :] = sf + else: + flat = scales.view(scales.shape[0], heads, tokens_per_block * groups) + offsets = quad_base[:, None] + _V_SF_QUAD * torch.arange(groups, device=sf.device) + for h in range(heads): + flat[page[:, None], h, offsets] = sf[:, h] + return tuple(out) + + +qsa_paged_kv_store_nvfp4 = QsaPagedKvStoreNvfp4() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa.py index c1af92df810c..d6ebc974a7d9 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa.py @@ -36,6 +36,7 @@ _SPLIT_K = ("splitk", "merge") _PREFILL = ("prefill",) +_FP8 = dict(cache_dtype=torch.float8_e4m3fn) def _spec(**overrides) -> dict: @@ -53,8 +54,9 @@ def _decode(segments, launches=_SPLIT_K, **overrides) -> dict: return _spec(**{**base, **overrides}) -def _mixed(requests, context_lens) -> dict: - return _spec(mixed=(requests, context_lens), slots=32, prefill_pair=True, launches=_PREFILL) +def _mixed(requests, context_lens, **overrides) -> dict: + base = dict(mixed=(requests, context_lens), slots=32, prefill_pair=True, launches=_PREFILL) + return _spec(**{**base, **overrides}) class QsaPagedSparseGqa(OpWrapper): @@ -118,6 +120,18 @@ class QsaPagedSparseGqa(OpWrapper): 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="decode MTP7: 1 request x 8 rows, 32-way split-K", + spec=_decode(((4000, 8),)), + ), + Cell( + why="decode MTP7: 8 requests x 8 rows, 8-way split-K", + spec=_decode(tuple((4000 + 37 * i, 8) for i in range(8))), + ), + Cell( + why="decode MTP7: 32 slots x 8 rows = 256 rows, 512 programs, 4-way split-K", + spec=_decode(tuple((100 + (499 * g) % 8000, 8) 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), @@ -185,6 +199,24 @@ class QsaPagedSparseGqa(OpWrapper): why="fp8-e4m3 pool under a bf16 q, read unscaled", spec=_decode(((900, 1),) * 4, cache_dtype=torch.float8_e4m3fn), ), + # The target's three call shapes over an fp8 pool (kv_cache_dtype=fp8, unit scale). + Cell( + why="fp8 decode MTP: 4 requests x 2 rows, split-K + merge", + spec=_decode(tuple((4000 + 37 * i, 2) for i in range(4)), **_FP8), + ), + Cell( + why="fp8 context 128 rows from 2032: pair passed, still split-K", + spec=_decode(((2032, 128),), prefill_pair=True, **_FP8), + ), + Cell( + why="fp8 context 272 rows from 2032: causal-prefix prefill kernel", + spec=_decode(((2032, 272),), launches=_PREFILL, prefill_pair=True, **_FP8), + ), + Cell(why="fp8 mixed 257 rows: 1 chunk + 31 MTP decodes", spec=_mixed(32, (195,), **_FP8)), + Cell( + why="fp8 decode MTP7: 32 slots x 8 rows = 256 rows, 512 programs, 4-way split-K", + spec=_decode(tuple((100 + (499 * g) % 8000, 8) for g in range(32)), **_FP8), + ), Cell( why="int64 request indices", spec=_decode(((1000, 2),) * 4, request_dtype=torch.int64) ), @@ -217,7 +249,8 @@ class docstring). QSA's producer -- `qsa_topk_select` with zero - 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. + - There is no KV scale argument: an fp8 pool is read at face value (each code + widens exactly to bf16 under a bf16 q), so it must hold K/V written at unit scale. - `D` must be a power of two (192 is refused); `Hq % Hkv == 0`. The op validates its own inputs loudly (11 `ValueError`s pinned by the @@ -229,6 +262,10 @@ class docstring). QSA's producer -- `qsa_topk_select` with zero launch for multi-split launches of 4 rows or fewer; both launch modes are inside the cells. Calls are bitwise deterministic. + The decode site carries up to 8 rows per request (max_draft_len 7): the + cells run 1, 8 and 32 requests of 8 rows, which is 16, 128 and 512 + `rows * Hkv` programs, so the 32-way, 8-way and 4-way split-K launches. + 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 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4.py new file mode 100644 index 000000000000..23f4f7bb3fee --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4.py @@ -0,0 +1,1156 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""QSA sparse GQA over an NVFP4 paged pool: attention over explicitly selected tokens.""" + +from __future__ import annotations + +import os +from typing import Optional + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._utils import get_sm_version + +#: Base-2 online softmax: scores are scaled by log2(e) once. +_LOG2_E = tl.constexpr(1.4426950408889634) +#: Values per e4m3 block scale, e2m1 values per byte, tokens per V scale group. +_SF_VEC = 16 +_PER_BYTE = 2 +_V_SF_QUAD = 4 +#: e2m1 codes 0..7 -> magnitude; bit 3 is the sign. +_E2M1_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) + +#: This checkpoint's geometry and the engine's QSA table. `global_k` / `global_v` are +#: `kv_scale_quant_orig[1]` / `[2]`; `slot_pages` role pages of one lifecycle slot hold +#: the call's planes, layer `layer`'s K and V among them (the trunk pool holds 24). +_CHECKPOINT = dict( + q_heads=24, + kv_heads=2, + head_dim=256, + tokens_per_block=32, + compress_ratio=4, + token_topk=2048, + q_dtype=torch.bfloat16, + # q is the first 24 heads of the packed [T, 24 + 2 + 2, 256] QKV buffer. + q_layout="qkv_slice", + request_dtype=torch.int32, + # The engine's table: 260 pages per slot, -1 past each sequence. + pages=260, + global_k=1.0, + global_v=1.0, + slot_pages=2, + layer=0, +) + +_SPLIT_K = ("splitk", "merge") +_PREFILL = ("prefill",) + + +def _spec(**overrides) -> dict: + """A cell spec: this checkpoint's geometry unless the cell says otherwise. + + `segments` is `(first position, rows)` per request, in row order, or `mixed` is + `(requests, context chunk lengths)`; `prefill_pair` says whether `query_positions` + / `compress_ratio` are passed, which the target does exactly when the step has a + context request; `pool_pages` fixes the pool's page count (the requests then sit + on its last pages); `capture` sweeps the decode site under a CUDA graph. + """ + return {**_CHECKPOINT, **overrides} + + +def _decode(segments, launches=_SPLIT_K, **overrides) -> dict: + base = dict(segments=segments, slots=32, prefill_pair=False, launches=launches) + return _spec(**{**base, **overrides}) + + +def _mixed(requests, context_lens, launches=_PREFILL, **overrides) -> dict: + base = dict(mixed=(requests, context_lens), slots=32, prefill_pair=True, launches=launches) + return _spec(**{**base, **overrides}) + + +def _short(segments, slots, pool_pages) -> dict: + """Decode rows at low positions over a small pool: one cell per class of the + plain-int `num_pages` and `num_requests` Triton specializes on.""" + return _decode(segments, slots=slots, pool_pages=pool_pages) + + +@triton.jit +def _e2m1_to_f32(code): + """e2m1 code (0..15, bit 3 the sign) -> its value, exactly.""" + code = code.to(tl.int32) + magnitude_code = code & 7 + exponent = magnitude_code >> 1 + mantissa = (magnitude_code & 1).to(tl.float32) + magnitude = tl.where( + exponent == 0, + 0.5 * mantissa, + (2.0 + mantissa) * (1 << exponent).to(tl.float32) * 0.25, + ) + return tl.where((code & 8) != 0, -magnitude, magnitude) + + +@triton.jit +def _load_keys( + k_cache, + k_scale_cache, + pages, + kv_head, + tokens, + valid, + dims, + k_global, + k_stride_page, + k_stride_head, + k_stride_token, + ks_stride_page, + ks_stride_head, + SCALE_GROUPS: tl.constexpr, +): + """K for the selected tokens as an fp32 [HEAD_DIM, BLOCK_N] tile: e2m1 x scale x global.""" + packed = tl.load( + k_cache + + pages[None, :] * k_stride_page + + kv_head * k_stride_head + + tokens[None, :] * k_stride_token + + dims[:, None] // 2, + mask=valid[None, :], + other=0, + ).to(tl.int32) + codes = (packed >> ((dims[:, None] % 2) * 4)) & 0xF + scales = ( + tl.load( + k_scale_cache + + pages[None, :] * ks_stride_page + + kv_head * ks_stride_head + + tokens[None, :] * SCALE_GROUPS + + dims[:, None] // 16, + mask=valid[None, :], + other=0, + ) + .to(tl.float8e4nv, bitcast=True) + .to(tl.float32) + ) + return _e2m1_to_f32(codes) * scales * k_global + + +@triton.jit +def _load_values( + v_cache, + v_scale_cache, + pages, + kv_head, + tokens, + valid, + dims, + v_global, + v_stride_page, + v_stride_head, + v_stride_token, + vs_stride_page, + vs_stride_head, + SCALE_GROUPS: tl.constexpr, +): + """V for the selected tokens as an fp32 [BLOCK_N, HEAD_DIM] tile; V's scales sit in + four-token groups, `[tpb / 4, D / 16, 4]` per head page.""" + packed = tl.load( + v_cache + + pages[:, None] * v_stride_page + + kv_head * v_stride_head + + tokens[:, None] * v_stride_token + + dims[None, :] // 2, + mask=valid[:, None], + other=0, + ).to(tl.int32) + codes = (packed >> ((dims[None, :] % 2) * 4)) & 0xF + scales = ( + tl.load( + v_scale_cache + + pages[:, None] * vs_stride_page + + kv_head * vs_stride_head + + (tokens[:, None] // 4) * (4 * SCALE_GROUPS) + + (dims[None, :] // 16) * 4 + + tokens[:, None] % 4, + mask=valid[:, None], + other=0, + ) + .to(tl.float8e4nv, bitcast=True) + .to(tl.float32) + ) + return _e2m1_to_f32(codes) * scales * v_global + + +@triton.jit +def _nvfp4_sparse_gqa_prefill_kernel( + q, + k_cache, + v_cache, + k_scale_cache, + v_scale_cache, + kv_scale_quant_orig, + block_table, + selected_tokens, + request_indices, + query_positions, + output, + softmax_scale, + num_cache_pages, + num_requests, + page_table_width, + q_stride_row, + q_stride_head, + k_stride_page, + k_stride_head, + k_stride_token, + v_stride_page, + v_stride_head, + v_stride_token, + ks_stride_page, + ks_stride_head, + vs_stride_page, + vs_stride_head, + block_table_stride_request, + selected_stride_row, + output_stride_row, + output_stride_head, + GROUP_SIZE: tl.constexpr, + TOKENS_PER_BLOCK: tl.constexpr, + TOPK: tl.constexpr, + BLOCK_TOPK: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + HEAD_DIM: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + kv_head = tl.program_id(1) + request = tl.load(request_indices + row).to(tl.int64) + valid_request = (request >= 0) & (request < num_requests) + safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1) + k_global = tl.load(kv_scale_quant_orig + 1) + v_global = tl.load(kv_scale_quant_orig + 2) + + head_offsets = tl.arange(0, BLOCK_M) + dim_offsets = tl.arange(0, HEAD_DIM) + query_heads = kv_head * GROUP_SIZE + head_offsets + query_values = tl.load( + q + row * q_stride_row + query_heads[:, None] * q_stride_head + dim_offsets[None, :], + mask=(head_offsets < GROUP_SIZE)[:, None], + other=0.0, + ) + + running_max = tl.full([BLOCK_M], -float("inf"), tl.float32) + running_sum = tl.zeros([BLOCK_M], tl.float32) + accumulator = tl.zeros([BLOCK_M, HEAD_DIM], tl.float32) + token_offsets = tl.arange(0, BLOCK_N) + visible_tokens = tl.load(query_positions + row).to(tl.int64) + 1 + selected_blocks = tl.minimum(visible_tokens // COMPRESS_RATIO, BLOCK_TOPK) + valid_tokens = selected_blocks * COMPRESS_RATIO + visible_tokens % COMPRESS_RATIO + valid_tiles = tl.cdiv(tl.minimum(valid_tokens, TOPK), BLOCK_N) + + for tile in tl.range(0, valid_tiles, num_stages=2): + selected_columns = tile * BLOCK_N + token_offsets + logical_tokens = tl.load( + selected_tokens + row * selected_stride_row + selected_columns, + mask=selected_columns < TOPK, + other=-1, + ) + safe_tokens = tl.maximum(logical_tokens, 0) + logical_pages = safe_tokens // TOKENS_PER_BLOCK + token_in_page = safe_tokens % TOKENS_PER_BLOCK + valid = ( + valid_request + & (selected_columns < TOPK) + & (logical_tokens >= 0) + & (logical_pages < page_table_width) + ) + physical_pages = tl.load( + block_table + + safe_request * block_table_stride_request + + tl.minimum(logical_pages, page_table_width - 1), + mask=valid, + other=-1, + ).to(tl.int64) + valid &= (physical_pages >= 0) & (physical_pages < num_cache_pages) + safe_pages = tl.maximum(physical_pages, 0) + + keys = _load_keys( + k_cache, + k_scale_cache, + safe_pages, + kv_head, + token_in_page, + valid, + dim_offsets, + k_global, + k_stride_page, + k_stride_head, + k_stride_token, + ks_stride_page, + ks_stride_head, + HEAD_DIM // 16, + ).to(query_values.dtype) + values = _load_values( + v_cache, + v_scale_cache, + safe_pages, + kv_head, + token_in_page, + valid, + dim_offsets, + v_global, + v_stride_page, + v_stride_head, + v_stride_token, + vs_stride_page, + vs_stride_head, + HEAD_DIM // 16, + ).to(query_values.dtype) + scores = tl.dot(query_values, keys) * softmax_scale * _LOG2_E + scores = tl.where(valid[None, :], scores, -float("inf")) + next_max = tl.maximum(running_max, tl.max(scores, axis=1)) + correction = tl.math.exp2(running_max - next_max) + probabilities = tl.where( + valid[None, :], + tl.math.exp2(scores - next_max[:, None]), + 0.0, + ) + accumulator = tl.dot( + probabilities.to(values.dtype), + values, + accumulator * correction[:, None], + ) + running_sum = running_sum * correction + tl.sum(probabilities, axis=1) + running_max = next_max + + normalized = tl.where( + running_sum[:, None] > 0, + accumulator / tl.maximum(running_sum[:, None], 1.0e-20), + 0.0, + ) + tl.store( + output + + row * output_stride_row + + query_heads[:, None] * output_stride_head + + dim_offsets[None, :], + normalized, + mask=(head_offsets < GROUP_SIZE)[:, None], + ) + + +@triton.jit +def _nvfp4_sparse_gqa_splitk_kernel( + q, + k_cache, + v_cache, + k_scale_cache, + v_scale_cache, + kv_scale_quant_orig, + block_table, + selected_tokens, + request_indices, + partial_output, + partial_lse, + output, + softmax_scale, + num_rows, + num_cache_pages, + num_requests, + page_table_width, + q_stride_row, + q_stride_head, + k_stride_page, + k_stride_head, + k_stride_token, + v_stride_page, + v_stride_head, + v_stride_token, + ks_stride_page, + ks_stride_head, + vs_stride_page, + vs_stride_head, + block_table_stride_request, + selected_stride_row, + output_stride_row, + output_stride_head, + GROUP_SIZE: tl.constexpr, + NUM_QUERY_HEADS: tl.constexpr, + TOKENS_PER_BLOCK: tl.constexpr, + TOPK: tl.constexpr, + NUM_SPLITS: tl.constexpr, + NUM_TILES: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + HEAD_DIM: tl.constexpr, + USE_PDL: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + kv_head = tl.program_id(1) + split = tl.program_id(2) + if USE_PDL: + tl.extra.cuda.gdc_wait() + request = tl.load(request_indices + row).to(tl.int64) + valid_request = (request >= 0) & (request < num_requests) + safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1) + k_global = tl.load(kv_scale_quant_orig + 1) + v_global = tl.load(kv_scale_quant_orig + 2) + + head_offsets = tl.arange(0, BLOCK_M) + dim_offsets = tl.arange(0, HEAD_DIM) + token_offsets = tl.arange(0, BLOCK_N) + query_heads = kv_head * GROUP_SIZE + head_offsets + query_values = tl.load( + q + row * q_stride_row + query_heads[:, None] * q_stride_head + dim_offsets[None, :], + mask=(head_offsets < GROUP_SIZE)[:, None], + other=0.0, + ) + + running_max = tl.full([BLOCK_M], -float("inf"), tl.float32) + running_sum = tl.zeros([BLOCK_M], tl.float32) + accumulator = tl.zeros([BLOCK_M, HEAD_DIM], tl.float32) + split_tile_start = split * NUM_TILES // NUM_SPLITS + split_tile_end = (split + 1) * NUM_TILES // NUM_SPLITS + + for tile in range(split_tile_start, split_tile_end): + selected_columns = tile * BLOCK_N + token_offsets + logical_tokens = tl.load( + selected_tokens + row * selected_stride_row + selected_columns, + mask=selected_columns < TOPK, + other=-1, + ) + safe_tokens = tl.maximum(logical_tokens, 0) + logical_pages = safe_tokens // TOKENS_PER_BLOCK + token_in_page = safe_tokens % TOKENS_PER_BLOCK + valid = ( + valid_request + & (selected_columns < TOPK) + & (logical_tokens >= 0) + & (logical_pages < page_table_width) + ) + physical_pages = tl.load( + block_table + + safe_request * block_table_stride_request + + tl.minimum(logical_pages, page_table_width - 1), + mask=valid, + other=-1, + ).to(tl.int64) + valid &= (physical_pages >= 0) & (physical_pages < num_cache_pages) + safe_pages = tl.maximum(physical_pages, 0) + + keys = _load_keys( + k_cache, + k_scale_cache, + safe_pages, + kv_head, + token_in_page, + valid, + dim_offsets, + k_global, + k_stride_page, + k_stride_head, + k_stride_token, + ks_stride_page, + ks_stride_head, + HEAD_DIM // 16, + ).to(query_values.dtype) + values = _load_values( + v_cache, + v_scale_cache, + safe_pages, + kv_head, + token_in_page, + valid, + dim_offsets, + v_global, + v_stride_page, + v_stride_head, + v_stride_token, + vs_stride_page, + vs_stride_head, + HEAD_DIM // 16, + ).to(query_values.dtype) + scores = tl.dot(query_values, keys) + scores *= softmax_scale * _LOG2_E + scores = tl.where(valid[None, :], scores, -float("inf")) + next_max = tl.maximum(running_max, tl.max(scores, axis=1)) + correction = tl.math.exp2(running_max - next_max) + probabilities = tl.where( + valid[None, :], + tl.math.exp2(scores - next_max[:, None]), + 0.0, + ) + accumulator = tl.dot( + probabilities.to(values.dtype), + values, + accumulator * correction[:, None], + ) + running_sum = running_sum * correction + tl.sum(probabilities, axis=1) + running_max = next_max + + has_values = running_sum > 0 + normalized = tl.where( + has_values[:, None], + accumulator / tl.maximum(running_sum[:, None], 1.0e-20), + 0.0, + ) + output_mask = (head_offsets < GROUP_SIZE)[:, None] + if NUM_SPLITS == 1: + tl.store( + output + + row * output_stride_row + + query_heads[:, None] * output_stride_head + + dim_offsets[None, :], + normalized, + mask=output_mask, + ) + else: + partial_lse_values = tl.where( + has_values, + running_max + tl.math.log2(tl.maximum(running_sum, 1.0e-20)), + -float("inf"), + ) + tl.store( + partial_output + + ((split * num_rows + row) * NUM_QUERY_HEADS + query_heads[:, None]) * HEAD_DIM + + dim_offsets[None, :], + normalized, + mask=output_mask, + ) + tl.store( + partial_lse + (split * num_rows + row) * NUM_QUERY_HEADS + query_heads, + partial_lse_values, + mask=head_offsets < GROUP_SIZE, + ) + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.jit +def _nvfp4_merge_splitk_kernel( + partial_output, + partial_lse, + output, + num_rows, + output_stride_row, + output_stride_head, + NUM_QUERY_HEADS: tl.constexpr, + NUM_SPLITS: tl.constexpr, + BLOCK_SPLITS: tl.constexpr, + HEAD_DIM: tl.constexpr, + USE_PDL: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + head = tl.program_id(1) + if USE_PDL: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + split_offsets = tl.arange(0, BLOCK_SPLITS) + dim_offsets = tl.arange(0, HEAD_DIM) + split_mask = split_offsets < NUM_SPLITS + lse = tl.load( + partial_lse + (split_offsets * num_rows + row) * NUM_QUERY_HEADS + head, + mask=split_mask, + other=-float("inf"), + ) + max_lse = tl.max(lse, axis=0) + has_values = max_lse > -float("inf") + weights = tl.math.exp2(tl.where(split_mask & has_values, lse - max_lse, -float("inf"))) + denominator = tl.sum(weights, axis=0) + partials = tl.load( + partial_output + + ((split_offsets[:, None] * num_rows + row) * NUM_QUERY_HEADS + head) * HEAD_DIM + + dim_offsets[None, :], + mask=split_mask[:, None], + other=0.0, + ) + merged = tl.sum(partials * weights[:, None], axis=0) + merged = tl.where(denominator > 0, merged / denominator, 0.0) + tl.store(output + row * output_stride_row + head * output_stride_head + dim_offsets, merged) + + +def _splitk_launch_config( + rows: int, num_kv_heads: int, selected_width: int +) -> tuple[int, int, int]: + """(BLOCK_N, splits, warps): the QSA split-K buckets by programs = rows x Hkv.""" + base_programs = rows * num_kv_heads + if base_programs <= 4: + block_n, target_splits, num_warps = 16, 64, 4 + elif base_programs < 32: + block_n, target_splits, num_warps = 16, 32, 4 + elif base_programs <= 256: + block_n, target_splits, num_warps = 64, 8, 2 + elif base_programs <= 512: + block_n, target_splits, num_warps = 64, 4, 2 + else: + block_n, target_splits, num_warps = 64, 1, 2 + num_tiles = triton.cdiv(selected_width, block_n) + max_useful_splits = 1 << (num_tiles.bit_length() - 1) + return block_n, min(max_useful_splits, target_splits), num_warps + + +def _pdl_enabled(rows: int) -> bool: + return rows <= 4 and os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" and get_sm_version() >= 90 + + +def _sparse_gqa_nvfp4( + *, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + k_scale_cache: torch.Tensor, + v_scale_cache: torch.Tensor, + kv_scale_quant_orig: torch.Tensor, + block_table: torch.Tensor, + selected_tokens: torch.Tensor, + request_indices: torch.Tensor, + tokens_per_block: int, + softmax_scale: float, + query_positions: Optional[torch.Tensor] = None, + compress_ratio: Optional[int] = None, +) -> torch.Tensor: + """Sparse GQA over the selected tokens of an NVFP4 HND paged pool (see the entry).""" + if q.ndim != 3: + raise ValueError("NVFP4 sparse GQA: q must be [rows, query_heads, head_dim]") + rows, num_q_heads, head_dim = q.shape + if q.dtype not in (torch.bfloat16, torch.float16) or q.stride(2) != 1: + raise ValueError("NVFP4 sparse GQA: q must be BF16 or FP16 with a unit last-dim stride") + if head_dim < 32 or head_dim & (head_dim - 1): + raise ValueError("NVFP4 sparse GQA: head_dim must be a power of two of at least 32") + if k_cache.ndim != 4 or v_cache.shape != k_cache.shape: + raise ValueError("NVFP4 sparse GQA: k_cache and v_cache must be matching 4-D planes") + num_pages, num_kv_heads = k_cache.shape[0], k_cache.shape[1] + if num_kv_heads <= 0 or num_q_heads % num_kv_heads: + raise ValueError("NVFP4 sparse GQA: query heads must be a multiple of KV heads") + if tokens_per_block <= 0 or tokens_per_block % _V_SF_QUAD: + raise ValueError("NVFP4 sparse GQA: tokens_per_block must be a positive multiple of 4") + data_shape = (num_pages, num_kv_heads, tokens_per_block, head_dim // _PER_BYTE) + scale_shape = (num_pages, num_kv_heads, tokens_per_block, head_dim // _SF_VEC) + for name, plane, shape in ( + ("k_cache", k_cache, data_shape), + ("v_cache", v_cache, data_shape), + ("k_scale_cache", k_scale_cache, scale_shape), + ("v_scale_cache", v_scale_cache, scale_shape), + ): + if plane.dtype != torch.uint8 or tuple(plane.shape) != shape or plane.stride(3) != 1: + raise ValueError( + f"NVFP4 sparse GQA: {name} must be uint8 {shape} with a unit last-dim stride, " + f"got {plane.dtype} {tuple(plane.shape)}" + ) + for name, plane in (("k_scale_cache", k_scale_cache), ("v_scale_cache", v_scale_cache)): + if plane.stride(2) != head_dim // _SF_VEC: + raise ValueError(f"NVFP4 sparse GQA: each head page of {name} must be contiguous") + if ( + kv_scale_quant_orig.dtype != torch.float32 + or kv_scale_quant_orig.shape != (3,) + or not kv_scale_quant_orig.is_contiguous() + ): + raise ValueError("NVFP4 sparse GQA: kv_scale_quant_orig must be a contiguous fp32 [3]") + if block_table.ndim != 2 or block_table.dtype != torch.int32 or block_table.stride(1) != 1: + raise ValueError("NVFP4 sparse GQA: block_table must be a row-contiguous 2-D int32 tensor") + if block_table.shape[0] <= 0 or block_table.shape[1] <= 0 or num_pages <= 0: + raise ValueError("NVFP4 sparse GQA: the pool and the table must be non-empty") + if ( + selected_tokens.ndim != 2 + or selected_tokens.dtype != torch.int32 + or selected_tokens.shape[0] != rows + or selected_tokens.shape[1] <= 0 + or selected_tokens.stride(1) != 1 + ): + raise ValueError( + "NVFP4 sparse GQA: selected_tokens must be a row-contiguous int32 [rows, width]" + ) + if ( + request_indices.shape != (rows,) + or request_indices.dtype not in (torch.int32, torch.int64) + or not request_indices.is_contiguous() + ): + raise ValueError("NVFP4 sparse GQA: request_indices must be one contiguous int per row") + if (query_positions is None) != (compress_ratio is None): + raise ValueError( + "NVFP4 sparse GQA: pass both query_positions and compress_ratio, or neither" + ) + if query_positions is not None and ( + query_positions.shape != (rows,) + or query_positions.dtype not in (torch.int32, torch.int64) + or not query_positions.is_contiguous() + or compress_ratio <= 0 + ): + raise ValueError( + "NVFP4 sparse GQA: query_positions must be one contiguous int per row and " + "compress_ratio positive" + ) + if not (softmax_scale > 0 and softmax_scale < float("inf")): + raise ValueError("NVFP4 sparse GQA: softmax_scale must be finite and positive") + operands = (q, k_cache, v_cache, k_scale_cache, v_scale_cache, kv_scale_quant_orig) + operands += (block_table, selected_tokens, request_indices) + operands += () if query_positions is None else (query_positions,) + if not all(t.is_cuda and t.device == q.device for t in operands): + raise ValueError("NVFP4 sparse GQA: every operand must be on q's CUDA device") + + group_size = num_q_heads // num_kv_heads + block_m = max(16, triton.next_power_of_2(group_size)) + output = torch.empty_like(q, memory_format=torch.contiguous_format) + if rows == 0: + return output + plane_strides = ( + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(2), + v_cache.stride(0), + v_cache.stride(1), + v_cache.stride(2), + k_scale_cache.stride(0), + k_scale_cache.stride(1), + v_scale_cache.stride(0), + v_scale_cache.stride(1), + ) + if query_positions is not None and rows * num_kv_heads > 512: + final_topk = selected_tokens.shape[1] + token_topk = final_topk - compress_ratio + 1 + if token_topk <= 0 or token_topk % compress_ratio: + raise ValueError( + "NVFP4 sparse GQA: the selection width does not come from this compression ratio" + ) + _nvfp4_sparse_gqa_prefill_kernel[(rows, num_kv_heads)]( + q, + k_cache, + v_cache, + k_scale_cache, + v_scale_cache, + kv_scale_quant_orig, + block_table, + selected_tokens, + request_indices, + query_positions, + output, + softmax_scale, + num_pages, + block_table.shape[0], + block_table.shape[1], + q.stride(0), + q.stride(1), + *plane_strides, + block_table.stride(0), + selected_tokens.stride(0), + output.stride(0), + output.stride(1), + GROUP_SIZE=group_size, + TOKENS_PER_BLOCK=tokens_per_block, + TOPK=final_topk, + BLOCK_TOPK=token_topk // compress_ratio, + COMPRESS_RATIO=compress_ratio, + BLOCK_M=block_m, + BLOCK_N=64, + HEAD_DIM=head_dim, + num_warps=2, + num_stages=2, + ) + return output + + block_n, num_splits, num_warps = _splitk_launch_config( + rows, num_kv_heads, selected_tokens.shape[1] + ) + num_tiles = triton.cdiv(selected_tokens.shape[1], block_n) + use_pdl = num_splits > 1 and _pdl_enabled(rows) + if num_splits == 1: + partial_output = partial_lse = output + else: + partial_output = torch.empty((num_splits, *q.shape), dtype=torch.float32, device=q.device) + partial_lse = torch.empty( + (num_splits, rows, num_q_heads), dtype=torch.float32, device=q.device + ) + _nvfp4_sparse_gqa_splitk_kernel[(rows, num_kv_heads, num_splits)]( + q, + k_cache, + v_cache, + k_scale_cache, + v_scale_cache, + kv_scale_quant_orig, + block_table, + selected_tokens, + request_indices, + partial_output, + partial_lse, + output, + softmax_scale, + rows, + num_pages, + block_table.shape[0], + block_table.shape[1], + q.stride(0), + q.stride(1), + *plane_strides, + block_table.stride(0), + selected_tokens.stride(0), + output.stride(0), + output.stride(1), + GROUP_SIZE=group_size, + NUM_QUERY_HEADS=num_q_heads, + TOKENS_PER_BLOCK=tokens_per_block, + TOPK=selected_tokens.shape[1], + NUM_SPLITS=num_splits, + NUM_TILES=num_tiles, + BLOCK_M=block_m, + BLOCK_N=block_n, + HEAD_DIM=head_dim, + USE_PDL=use_pdl, + num_warps=num_warps, + num_stages=2, + launch_pdl=use_pdl, + ) + if num_splits == 1: + return output + _nvfp4_merge_splitk_kernel[(rows, num_q_heads)]( + partial_output, + partial_lse, + output, + rows, + output.stride(0), + output.stride(1), + NUM_QUERY_HEADS=num_q_heads, + NUM_SPLITS=num_splits, + BLOCK_SPLITS=triton.next_power_of_2(num_splits), + HEAD_DIM=head_dim, + USE_PDL=use_pdl, + num_warps=2, + num_stages=1, + launch_pdl=use_pdl, + ) + return output + + +def _dequantize_tokens(data, scales, pages, slots, global_scale, grouped: bool) -> torch.Tensor: + """[..., Hkv, D] fp32 values of the tokens at (pages, slots) of one NVFP4 role plane.""" + heads, tokens_per_block, packed = data.shape[1], data.shape[2], data.shape[3] + head_dim = packed * _PER_BYTE + groups = head_dim // _SF_VEC + raw = data[pages, :, slots, :].long() # [..., H, D/2] + lut = torch.tensor(_E2M1_VALUES, dtype=torch.float32, device=data.device) + lut = torch.cat([lut, -lut]) + values = torch.stack([lut[raw & 0xF], lut[raw >> 4]], dim=-1).flatten(-2) + if grouped: + flat = scales.reshape(scales.shape[0], heads, tokens_per_block * groups) + base = (slots // _V_SF_QUAD) * (_V_SF_QUAD * groups) + slots % _V_SF_QUAD + columns = base[..., None] + _V_SF_QUAD * torch.arange(groups, device=data.device) + sf = torch.stack([flat[pages[..., None], h, columns] for h in range(heads)], dim=-2) + else: + sf = scales[pages, :, slots, :] + sf = sf.view(torch.float8_e4m3fn).float().repeat_interleave(_SF_VEC, dim=-1) + return values * sf * global_scale + + +class QsaPagedSparseGqaNvfp4(OpWrapper): + """Sparse GQA over the selected tokens of an NVFP4 HND paged pool. + + Triton kernels defined in this file. For row `r`, query head `h`, KV head + `h // (Hq // Hkv)` (contiguous groups), over the columns `j` of + `selected_tokens[r]` that name a resident token:: + + t = selected_tokens[r, j] # logical position + page = block_table[request_indices[r], t // tokens_per_block] + K[t] = bf16(e2m1(k_cache) * e4m3(k_scale_cache) * kv_scale_quant_orig[1]) + V[t] = bf16(e2m1(v_cache) * e4m3(v_scale_cache) * kv_scale_quant_orig[2]) + s_j = (q[r, h] . K[t]) * softmax_scale + out[r, h] = sum_j softmax(s)_j * V[t] + + The pool layout is the attention op's own NVFP4 append: e2m1 nibbles per + head page `[tpb, D/2]` (low nibble = even element), K's e4m3 scales + token-major `[tpb, D/16]`, V's in four-token groups `[tpb/4, D/16, 4]`. + Each value is dequantized in fp32 and rounded once to `q`'s dtype; at unit + second-level scales every NVFP4 value is exact in bf16, so a call is the + bf16 sparse GQA over the bf16 values the pool holds. After the dequantize the + kernels are QSA's bf16 sparse GQA kernels operation for operation: the same + tiles, online softmax and split-K merge. + + A negative index, a page past the table's width, a negative or out-of-pool + table entry, and a request index outside the table's rows contribute + nothing; a row left with nothing returns zeros. Returns a fresh, contiguous + `[rows, Hq, D]` tensor of `q.dtype`. + + Two kernels, chosen inside the call: with `query_positions` and + `compress_ratio` given and `rows * Hkv > 512`, a causal-prefix prefill + kernel; otherwise split-K, plus a merge kernel when it used more than one + split. **The prefill pair is a bound, not a mask**, exactly as on the bf16 + entry: the prefill kernel reads only `min(visible // cr, block_topk) * cr + + visible % cr` leading columns of a row. + + Fusion boundary: the block-table gather, the NVFP4 dequantize, scale, + online softmax, P.V and the split-K merge are inside. Selection, appending + K/V to the pool, and q's projection, norm, RoPE and output gate are not. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode: 1 row past the budget, 64-way 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 draft 3: 4 requests x 4 rows", + spec=_decode(tuple((4000 + 37 * i, 4) for i in range(4))), + ), + Cell( + why="decode: 15 rows (5 requests x 3), 32-way split-K", + spec=_decode(tuple((4000 + 37 * i, 3) for i in range(5))), + ), + Cell( + why="decode MTP draft 7: 16 requests x 8 rows (128 rows, 8-way)", + spec=_decode(tuple((100 + (499 * g) % 8000, 8) for g in range(16))), + ), + Cell( + why="decode MTP draft 7: 32 x 8 = 256 rows, 512 programs (4-way top)", + spec=_decode(tuple((100 + (499 * g) % 8000, 8) 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), + ), + Cell( + why="context 2 rows with the pair: 64-way split-K, PDL", + spec=_decode(((2032, 2),), prefill_pair=True), + ), + 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: the 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 4096 rows from 0: rows see 1..4096 tokens", + spec=_decode(((0, 4096),), launches=_PREFILL, prefill_pair=True), + ), + Cell(why="mixed 257 rows: 1 chunk + 31 MTP decodes", spec=_mixed(32, (195,))), + Cell( + why="mixed 3806 rows: 9 chunks + 23 MTP decodes", + spec=_mixed(32, (1024, 900, 700, 500, 300, 200, 100, 30, 6)), + ), + Cell( + why="second-level scales g_k 0.37, g_v 2.5: the dequantize rounds to bf16", + spec=_decode(tuple((4000 + 37 * i, 2) for i in range(4)), global_k=0.37, global_v=2.5), + ), + Cell(why="fp16 q", spec=_decode(((900, 1),) * 4, q_dtype=torch.float16)), + Cell( + why="int64 request indices", spec=_decode(((1000, 2),) * 4, request_dtype=torch.int64) + ), + Cell(why="contiguous q", spec=_decode(((900, 1),) * 4, q_layout="contiguous")), + Cell( + why="holes: all-padding row, foreign request, unmapped page", + spec=_decode(((900, 1),) * 4, holes=True), + ), + Cell( + why="the engine's trunk layout: layer 11's planes among the 24 role pages of a slot", + spec=_decode( + tuple((4000 + 37 * i, 4) for i in range(4)), slot_pages=24, layer=11, pool_pages=600 + ), + ), + Cell( + why="the trunk layout, context 272 rows from 2032: the prefill kernel", + spec=_decode( + ((2032, 272),), + launches=_PREFILL, + prefill_pair=True, + slot_pages=24, + layer=11, + pool_pages=600, + ), + ), + # Triton compiles once per class of each plain-int argument -- equal to 1, a multiple + # of 16, or anything else: the pool's page count and the table's height, and the + # split-K kernel's row count (rows 1, 15, 16, 128 and 256 above). + Cell( + why="specialization: pool of 100 pages (other), table height 1", + spec=_short(((600, 1),), slots=1, pool_pages=100), + ), + Cell( + why="specialization: pool of 100 pages (other), table height 8", + spec=_short(((600, 1),) * 4, slots=8, pool_pages=100), + ), + Cell( + why="specialization: pool of 100 pages (other), table height 16", + spec=_short(((600, 1),) * 4, slots=16, pool_pages=100), + ), + Cell( + why="specialization: pool of 96 pages (multiple of 16), table height 1", + spec=_short(((600, 1),), slots=1, pool_pages=96), + ), + Cell( + why="specialization: pool of 96 pages (multiple of 16), table height 8", + spec=_short(((600, 1),) * 4, slots=8, pool_pages=96), + ), + Cell( + why="specialization: pool of 96 pages (multiple of 16), table height 16", + spec=_short(((600, 1),) * 4, slots=16, pool_pages=96), + ), + Cell( + why="captured decode: G 1, 8 and 32 requests x 1 and 8 rows, replay == eager", + spec=_decode((), slots=32, capture=dict(groups=(1, 8, 32), rows_per_request=(1, 8))), + ), + ) + + note = """ + Metadata consumed: none -- every input is an argument. What the target + passes: `block_table` is `qsa_block_table` (physical page ids, the unscaled + table); `selected_tokens` is `qsa_expand_block_indices`' output, logical + positions, `-1`-padded; `k_cache` / `v_cache` and `k_scale_cache` / + `v_scale_cache` are the layer's planes of the QSA manager's `get_buffers(i, + "HND")` and `get_block_scale_buffers(i, "HND")`; `kv_scale_quant_orig` is the + fp32 [3] tensor its dense attention call binds (index 1 for K, 2 for V). + + Call sites (the target's `modeling.py`): the context/mixed forward passes the + prefill pair on every step, the decode forward never does; both pass `q` as + the 24-head slice of the packed QKV row (row stride 7168). + + Launch shapes at Hkv 2 (programs = rows x 2), as on the bf16 entry: split-K + 64-way up to 2 rows, 32-way up to 15 (16-token tiles), 8-way up to 128 and + 4-way up to 256 (64-token tiles), each followed by the merge; past 256 rows + the prefill kernel when the pair is passed, else one split and no merge. + The decode site reaches 256 rows at max_draft_len 7 (32 requests x 8). + + Equivalence: the test also runs QSA's bf16 sparse GQA over a bf16 pool + holding the dequantized values; on every cell the two outputs agree within + the test's band. They are not bitwise equal: the dequantizing load changes + the compiled kernel around the same arithmetic. + + CUDA graphs: the decode site captures at G requests x 1..8 rows over the + [32, 260] table; a replay after new operands are written in place is + bitwise the eager call on them. + + Compiled variants: the kernels take the pool's page count (`num_pages`), the + table's row count (`num_requests`) and the split-K kernel the step's row + count as plain ints, and Triton compiles each apart at 1, at multiples of 16 + and otherwise. Every other plain int -- the strides, the table width 260 and + the selection's row stride 2051 -- is fixed at the served value the cells use. + + Preconditions: + - Whenever the prefill kernel can run, each selection row is packed at the + front with exactly the bound's count of valid entries (QSA's producer + does this; see the bf16 entry's note). + - The pool holds what the attention op's NVFP4 append writes; the + `qsa_paged_kv_store_nvfp4` entry writes the same bits. + - `D` a power of two >= 32; `Hq % Hkv == 0`; `tokens_per_block % 4 == 0`. + + The callable rejects loudly (ValueError) every operand whose dtype, shape or + layout it cannot read; `rows == 0` returns an empty `[0, Hq, D]` tensor. + Calls are bitwise deterministic. + + Every row, request, page and position index is widened to 64 bits before any + product. Not certified: head geometries other than 24/2 x 256; + `tokens_per_block` other than 32; more than 8192 rows; capture of the + context/mixed site; a pool of one slot. + """ + + def __call__( + self, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + k_scale_cache: torch.Tensor, + v_scale_cache: torch.Tensor, + kv_scale_quant_orig: torch.Tensor, + block_table: torch.Tensor, + selected_tokens: torch.Tensor, + request_indices: torch.Tensor, + tokens_per_block: int, + softmax_scale: float, + query_positions: Optional[torch.Tensor] = None, + compress_ratio: Optional[int] = None, + ) -> torch.Tensor: + return _sparse_gqa_nvfp4( + q=q, + k_cache=k_cache, + v_cache=v_cache, + k_scale_cache=k_scale_cache, + v_scale_cache=v_scale_cache, + kv_scale_quant_orig=kv_scale_quant_orig, + block_table=block_table, + selected_tokens=selected_tokens, + request_indices=request_indices, + tokens_per_block=tokens_per_block, + softmax_scale=softmax_scale, + query_positions=query_positions, + compress_ratio=compress_ratio, + ) + + def reference( + self, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + k_scale_cache: torch.Tensor, + v_scale_cache: torch.Tensor, + kv_scale_quant_orig: torch.Tensor, + block_table: torch.Tensor, + selected_tokens: torch.Tensor, + request_indices: torch.Tensor, + tokens_per_block: int, + softmax_scale: float, + query_positions: Optional[torch.Tensor] = None, + compress_ratio: Optional[int] = None, + ) -> torch.Tensor: + """fp32 sparse GQA over every resident selected token, rounded once. + + K and V are dequantized in fp32 and rounded to `q`'s dtype, as the kernel + does. `query_positions` and `compress_ratio` are ignored: on a + prefix-packed selection the prefill kernel's bound reads exactly the valid + entries. Walked 32 rows at a time to bound the gathered K/V. + """ + rows, num_q_heads, head_dim = q.shape + num_pages, num_kv_heads = k_cache.shape[0], k_cache.shape[1] + groups = num_q_heads // num_kv_heads + num_requests, table_width = block_table.shape + requests = request_indices.long() + out = torch.empty(rows, num_q_heads, head_dim, dtype=torch.float32, device=q.device) + for start in range(0, rows, 32): + stop = min(start + 32, rows) + chunk = selected_tokens[start:stop].long() + logical_page = chunk.clamp_min(0) // tokens_per_block + in_page = chunk.clamp_min(0) % tokens_per_block + request = requests[start:stop] + owned = (request >= 0) & (request < num_requests) + physical = block_table[ + request.clamp(0, num_requests - 1)[:, None].expand_as(logical_page), + logical_page.clamp(max=table_width - 1), + ].long() + resident = ( + (chunk >= 0) + & owned[:, None] + & (logical_page < table_width) + & (physical >= 0) + & (physical < num_pages) + ) + physical = physical.clamp(0, num_pages - 1) + keys = _dequantize_tokens( + k_cache, k_scale_cache, physical, in_page, kv_scale_quant_orig[1], grouped=False + ) + values = _dequantize_tokens( + v_cache, v_scale_cache, physical, in_page, kv_scale_quant_orig[2], grouped=True + ) + keys = keys.to(q.dtype).float() + values = values.to(q.dtype).float() + grouped = q[start:stop].reshape(-1, num_kv_heads, groups, head_dim).float() + scores = torch.einsum("bhgd,bkhd->bhgk", grouped, keys) * softmax_scale + scores = scores.masked_fill(~resident[:, None, None, :], float("-inf")) + weights = torch.nan_to_num(torch.softmax(scores, dim=-1), nan=0.0) + pooled = torch.einsum("bhgk,bkhd->bhgd", weights, values) + out[start:stop] = pooled.reshape(stop - start, num_q_heads, head_dim) + return out.to(q.dtype) + + +qsa_paged_sparse_gqa_nvfp4 = QsaPagedSparseGqaNvfp4() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_prefill_compress.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_prefill_compress.py index ced33a120477..8efd444d874a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_prefill_compress.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_prefill_compress.py @@ -38,6 +38,10 @@ _DECODE = ((4095, 1), (4096, 1), (4097, 1), (4098, 1), (4099, 1), (8191, 1), (3, 1), (2047, 1)) +#: 32 requests of one MTP7 verify block each (8 rows, two completed groups), at starts that put +#: some blocks across a page edge and some inside one. +_MTP7_32 = tuple((29 + 61 * i, 8) for i in range(32)) + def _spec(segments, pages, **overrides) -> dict: """`segments` is `(first position, rows)` per request, in row order.""" @@ -114,6 +118,14 @@ class QsaPrefillCompress(OpWrapper): 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="decode MTP7: 1 request x 8 rows, two anchors, None", + spec=_spec(((4090, 8),), pages=160), + ), + Cell( + why="decode MTP7: 32 requests x 8 rows, two anchors each, None", + spec=_spec(_MTP7_32, pages=1100, slots=40, rounding_flips=1), + ), Cell( why="mixed: chunk from 0, chunk from 1000, 2 MTP decodes, None", spec=_spec(((0, 64), (1000, 36), (4094, 2), (255, 2)), pages=256), @@ -142,6 +154,18 @@ class QsaPrefillCompress(OpWrapper): why="index cache with a head-dim stride of 2", spec=_old(((0, 64),), pages=2, layout="head_dim_strided"), ), + # Under an NVFP4 KV pool the position page shares each physical page with the index + # pages, so both caches are strided views of one allocation: index page stride + # 12 x 4096 + 192 = 49344 on the trunk (4288 on the draft), position strides + # (24672, 3, 1) (2144, 3, 1). + Cell( + why="NVFP4 trunk planes: 12 index roles and the position page share each page", + spec=_old(((0, 64),) * 2, pages=4, layout="shared_page", roles=12), + ), + Cell( + why="NVFP4 draft planes: one index role and the position page share each page", + spec=_old(((0, 64),) * 2, pages=4, layout="shared_page", roles=1), + ), Cell( why="every row twice: two rows per anchor in one launch", spec=_old(((0, 64),) * 2, pages=4, duplicate=True), @@ -180,7 +204,17 @@ class QsaPrefillCompress(OpWrapper): `cos_sin`'s dtype is the rotation's precision: bf16 is accepted and was measured 1.4% of elements off by up to 1.56e-2 -- not certified. Text rows reduce mRoPE to plain RoPE bit for bit. Plain `@triton.jit`, no - autotune; CUDA-graph capturable at a fixed row count. + autotune; CUDA-graph capturable at a fixed row count (captured at 128 rows + and at 256, 32 MTP7 verify blocks). + + Under an NVFP4 KV pool both caches are strided views of one per-slot + allocation (each physical page holds the index roles, then the position + page); the cells check the other index roles and the position page + untouched. + + The kernel's `tl.rsqrt` is the hardware approximation, an fp32 ulp from the + reference's: every cell is bitwise the reference but the 32-block MTP7 one, where + 1 of 8192 anchor elements rounds one bf16 step away. Hard limits (read from the kernel, not probed): positions, requests and page ids are widened to int64 before any address arithmetic, so no int32 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_select.py index 1903a9cc6f32..fda108df6ac2 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_select.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_select.py @@ -100,6 +100,31 @@ class QsaTopkSelect(OpWrapper): tail="nan", ), ), + Cell( + why="decode MTP7: one request's 8 rows (6000..6007), NaN tail", + spec=dict( + rows=8, + cols=_SCORE_COLUMNS, + index_topk=_BLOCK_TOPK, + windows=("positions", tuple(range(6000, 6008))), + logits="relu", + tail="nan", + ), + ), + Cell( + why="decode MTP7: 32 requests x 8 rows = 256 rows, some below the budget, NaN tail", + spec=dict( + rows=256, + cols=_SCORE_COLUMNS, + index_topk=_BLOCK_TOPK, + windows=( + "positions", + tuple(600 + 233 * g + j for g in range(32) for j in range(8)), + ), + logits="relu", + tail="nan", + ), + ), Cell( why="context: 8191 rows at positions 0..8190 (warm-up capture)", spec=dict( @@ -222,9 +247,10 @@ class QsaTopkSelect(OpWrapper): nothing. `torch.ops.trtllm.indexer_topk_decode` is the request-grouped sibling; QSA does not use it. - CUDA graphs: captured at `[8, 2048]`, `k=512`, and replayed against new - logits and a new `row_ends` through the captured addresses. Bitwise - comparison of two replays is not a valid check (order is unspecified). + CUDA graphs: captured at `[8, 2048]` and `[256, 2048]` (32 requests at + max_draft_len 7), `k=512`, and replayed against new logits and a new + `row_ends` through the captured addresses. Bitwise comparison of two + replays is not a valid check (order is unspecified). Hard limits: not probed. `indices` is addressed as `row * index_topk + slot`; the largest certified extent is 12416 x 64. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/vision_rope_2d.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/vision_rope_2d.py index 559ba69efc9a..7a4388ccbb9b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/vision_rope_2d.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/vision_rope_2d.py @@ -140,6 +140,19 @@ class VisionRope2d(OpWrapper): grid_dtype=torch.int64, ), ), + Cell( + why="call site: packed qkv views, video at the 32768-row ceiling", + spec=dict( + grid=((64, 16, 32),), + dtype=torch.bfloat16, + layout="packed_qkv", + grid_dtype=torch.int64, + ), + ), + Cell( + why="float32, video at the 32768-row ceiling, two frame groups of 128x128", + spec=dict(grid=((2, 128, 128),), dtype=torch.float32), + ), Cell( why="grid_thw int64 on the device (host sync)", spec=dict( @@ -177,7 +190,7 @@ class VisionRope2d(OpWrapper): measured. Not certified: mixed q/k dtypes; float8; `spatial_merge_size` != 2; grid - sides past 6272 or calls past 22144 tokens. + sides past 6272 or calls past 32768 rows. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py new file mode 100644 index 000000000000..a5bff838b059 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py @@ -0,0 +1,535 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CuteDSL NVFP4 FC1 of the routed experts: gather, grouped GEMM, SwiGLU and NVFP4 +requantization of the intermediate, in one trtllm call.""" + +from __future__ import annotations + +from typing import NamedTuple, Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import SWIGLU_LIMIT_SCALAR_DISABLED +from tensorrt_llm._torch.utils import ActivationType + +_SWIGLU = int(ActivationType.Swiglu) + +#: e2m1 magnitudes by 3-bit code, and the midpoints between neighbours. A value +#: exactly on a midpoint rounds to the even code (`_TIE_UP[i]`: the even side is up). +E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) +E2M1_MAX = 6.0 +E4M3_MAX = 448.0 +_E2M1_MIDPOINTS = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0) +_TIE_UP = (False, True, False, True, False, True, False) + +#: The pairing the fused SwiGLU epilogue reads: the weight's rows `[128 g, 128 g + 64)` +#: are up rows, the next 64 their gate rows, and output column `64 g + r` pairs row +#: `128 g + r` with row `128 g + 64 + r`. Fixed by the epilogue's 64-column subtile, +#: whatever MMA N tile the autotuner picks. +GATE_UP_GROUP = 64 + + +def _ckpt( + tokens: int, + tile_size: int, + mma_n: int = 128, + dist: str = "router", + gsf: str = "calibrated", + zero_tokens: int = 0, +) -> dict: + """The checkpoint's routed experts: 512 experts, top-10, hidden 2560, I 640, all + local (tp1), SwiGLU, no clamp, 16-element e4m3 block scales. + + `mma_n` is the autotuner tactic's MMA N tile (128 is also the fallback when the + tuner has no entry). `dist` names how the test draws the expert ids, as in the + moe_sort entry's test: `router` = each row's top-10 of uniform scores; `hot` = + every row on the same 10 experts; `tile_edges` = three experts hold exactly + tile-1, tile and tile+1 rows; `round_robin` = row i takes experts (10 i + j) mod + 512. `gsf` sets `global_sf` against the activations' magnitude: `calibrated` = + the activation maximum the test calibrates to maps to the e4m3 maximum, as the + checkpoint's per-layer `1 / down_proj.input_scale` does; `saturating` = 25x + above that, so most block scales clip at 448 and the codes at 6; `subnormal` = + 2^-12 of it, so block scales fall into e4m3 subnormals and zero. `zero_tokens` + makes that many leading tokens all-zero rows (zero block scales). + """ + return dict( + tokens=tokens, + tile_size=tile_size, + mma_n=mma_n, + dist=dist, + gsf=gsf, + zero_tokens=zero_tokens, + num_experts=512, + top_k=10, + hidden=2560, + intermediate=640, + ) + + +class Fc1Expected(NamedTuple): + """What `reference` returns: the op's two outputs, plus what the test's compare reads. + + `c` / `c_sf` are laid out exactly as the op returns them; rows the op does not + write hold zeros. `rows` marks the permuted rows the op writes and the compare + gates. `scaled` is `y * global_sf` on those rows, before rounding -- the value + the block scale and the codes are derived from, which rounding loses. + """ + + c: torch.Tensor + c_sf: torch.Tensor + rows: torch.Tensor + scaled: torch.Tensor + + +class CuteDslNvfp4GatherGroupedGemmActFusion(OpWrapper): + """`torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell`: + FC1 of the routed experts on the CuteDSL MoE path, one kernel launch. + + For every permuted row `p` the sort filled -- tile `i = p // tile_size` below + `num_non_exiting_tiles` and `p < tile_idx_to_mn_limit[i]`: + + t = permuted_idx_to_expanded_idx[p] // top_k the token + e = tile_idx_to_group_idx[i] the local expert + acc = alpha[e] * deq(input[t]) @ deq(weight[e]).T [2 I], fp32 + y[p, j] = up * silu(gate), up = acc[128 g + r], gate = acc[128 g + 64 + r], + j = 64 g + r + out[p], out_sf[p] = NVFP4(y[p] * global_sf) 16-element blocks + + where `deq` is e2m1 times its e4m3 block scale, and NVFP4 takes per block + `sf = e4m3(max|v| / 6)` (saturating at 448) and codes `e2m1(v / sf)` (saturating + at 6). So `out * out_sf / global_sf` approximates `y`. + + Arguments, as the CuteDSL MoE passes them (`run_moe_nvfp4_impl`, + `_torch/moe/fused_moe/fused_moe_cute_dsl.py`): + + - `input` `[T, H/2]` float4_e2m1fn_x2 and `input_scale` `[T, H/16]` uint8: the + un-permuted activation as `fp4_quantize(x, 1 / input_scale_1, 16, False, + False)` returns it -- LINEAR (row-major) e4m3 scales, viewed 2-D. Gathered + inside; the caller never permutes it. + - `weight` `[E, 2I, H/2]` float4_e2m1fn_x2: per expert `cat([up, gate])` + interleaved in `GATE_UP_GROUP`-row groups (up rows 0:64, gate rows 0:64, up + rows 64:128, ...). No 32-row shuffle (that is trtllm-gen's layout). + - `weight_scale` `[E, 2I, H/16]` uint8: the e4m3 block scales with the same row + interleave, then 128x4-swizzled per expert (`block_scale_interleave`). + - `alpha` `[E]` fp32: per expert, `input_scale_1 * weight_scale_2[e]` -- one + factor for both halves, applied to the accumulator before SwiGLU. + - the four sort tensors, exactly as `moe_sort` returns them for the same + `top_k` and `tile_tokens_dim == tile_size`. + - `global_sf` fp32, one element: the FC2 input's global scale, `1 / + max(down_proj.input_scale)` of the layer. + - `num_experts`, `num_local_experts`, `local_expert_offset`: read only by the + autotuner's shape buckets and cache key; the kernel takes the expert count + from `weight.size(0)` and the expert from the tile. + + Returns `(c, c_sf)`: `c` `[M, I/2]` float4_e2m1fn_x2 and `c_sf` `[M * I/16]` + uint8, `M = permuted_idx_to_expanded_idx.numel()`, one row per PERMUTED row and + the block scales 128x4-swizzled over `[M, I/16]` -- exactly the `input` / + `input_scale` the plain NVFP4 grouped GEMM (FC2) takes, with the same four sort + tensors. Rows the sort did not fill are not part of the result (see `note`). + + Fusion boundary: the gather of the activation rows, the grouped GEMM, the + per-expert alpha, SwiGLU and the requantization happen inside. The caller owns + the activation quantization, the sort, FC2 and the routing-weighted combine; + routing weights are not read here. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = tuple( + cell + for tile in (128, 256) + for cell in ( + Cell( + why=f"tile {tile}: decode, 1 request without MTP", + spec=_ckpt(1, tile), + ), + Cell( + why=f"tile {tile}: 1 request at MTP verify width 2", + spec=_ckpt(2, tile), + ), + Cell( + why=f"tile {tile}: 1 request at MTP verify width 4", + spec=_ckpt(4, tile), + ), + Cell( + why=f"tile {tile}: decode, 8 requests", + spec=_ckpt(8, tile), + ), + Cell( + why=f"tile {tile}: decode at the 32-request ceiling", + spec=_ckpt(32, tile), + ), + Cell( + why=f"tile {tile}: 32 requests at MTP verify width 3", + spec=_ckpt(96, tile), + ), + Cell( + why=f"tile {tile}: 32 requests at MTP verify width 4; largest captured T", + spec=_ckpt(128, tile), + ), + Cell( + why=f"tile {tile}: 8192 rows, max_num_tokens prefill", + spec=_ckpt(8192, tile), + ), + Cell( + why=f"tile {tile}: 8192 rows all on 10 experts; {8192 // tile} full tiles each", + spec=_ckpt(8192, tile, dist="hot"), + ), + Cell( + why=f"tile {tile}: expert rows tile-1, tile, tile+1; the mn_limit edge", + spec=_ckpt(300, tile, dist="tile_edges"), + ), + Cell( + why=f"tile {tile}: 510 experts with one row each; every tile partial", + spec=_ckpt(51, tile, dist="round_robin"), + ), + Cell( + why=f"tile {tile}: MMA N 256 tactic, decode", + spec=_ckpt(1, tile, mma_n=256), + ), + Cell( + why=f"tile {tile}: MMA N 256 tactic, 128 rows", + spec=_ckpt(128, tile, mma_n=256), + ), + Cell( + why=f"tile {tile}: MMA N 256 tactic, 8192-row prefill", + spec=_ckpt(8192, tile, mma_n=256), + ), + Cell( + why=f"tile {tile}: global_sf 25x calibrated; block scales clip at 448", + spec=_ckpt(2048, tile, gsf="saturating"), + ), + Cell( + why=f"tile {tile}: tiny global_sf; subnormal and zero block scales", + spec=_ckpt(2048, tile, gsf="subnormal"), + ), + Cell( + why=f"tile {tile}: 7 all-zero token rows; zero block scales", + spec=_ckpt(64, tile, zero_tokens=7), + ), + ) + ) + + note = """ + Path certified: the CuTe DSL kernel `BlockScaledContiguousGatherGroupedGemmKernel` + (`_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py`), + built by `Sm100BlockScaledContiguousGatherGroupedGemmActFusionRunner` + (`_torch/custom_ops/cute_dsl_custom_ops.py`), which JIT-compiles it with + `cute.compile` on the first call per (tile, tactic, top_k, activation) and + keeps it for the process. The op is registered only when the CuTe DSL imports + (`IS_CUTLASS_DSL_AVAILABLE`), and the runner accepts SM 100 and 103 only. + + Metadata consumed: none beyond the arguments. Stateless: allocates `c` and + `c_sf` on every call, on the current stream, with no host synchronization, so + it captures into a CUDA graph together with moe_sort. + + Call sites: the engine's CuteDSL MoE (`run_moe_nvfp4_impl`) and the target's + routed-expert step, eager for prefill and mixed batches (T up to + max_num_tokens 8192) and captured for decode (T = requests x MTP verify width, + up to 32 x 4 = 128). Tile 128 is the target's; 256 is certified as well. + + Tactics. The op asks the AutoTuner for an MMA tile of (tile_size, 128) or + (tile_size, 256) and with no tuning entry falls back to (tile_size, 128); + serving warmup tunes per row bucket and may pick either. Every pair's codes + and block scales are bitwise identical across tile 128 / 256 x MMA N 128 / 256 + x the fallback, so cells driven at one tactic certify the others. + + Determinism. A (token, expert) pair's output row depends on that pair alone: + it is bitwise unchanged by reordering an expert's rows (moe_sort orders them by + atomics from T = 5), by repeating the call, and by cutting a batch to a prefix. + + What is written: the rows the sort filled hold the result; the padding rows of + a written tile -- from its tile_idx_to_mn_limit to its end -- get zero block + scales and zero codes, so a plain grouped GEMM over whole tiles (FC2) reads + zeros there; the tiles at and past num_non_exiting_tiles are not written at + all, and under a graph keep whatever the buffer last held. FC2 must be tiled by + the same tile arrays, and the combine reads only + expanded_idx_to_permuted_idx's rows. + + Scale range: with global_sf 25x the calibrated value the block scales clip at + 448 and the codes at 6 (no NaN); with a tiny one the block scales fall to e4m3 + subnormals and to zero. An all-zero token gives scale 0 and codes of magnitude + 0. A block whose scale is 0 dequantizes to zero whatever its codes hold, and + the test's compare does not read them. + + Preconditions the caller owns (each wrong preparation ran without error, with + finite scales, and put most block scales and codes outside the test's gate): + - `weight` / `weight_scale` in the layout the class docstring states: not + trtllm-gen's row interleave, up rows first, scales swizzled after the same + row interleave as the weight, alpha not shifted by an expert. + - One `alpha` per expert scales BOTH halves, so up_proj and gate_proj must + share `weight_scale_2`, and `input_scale` must be one value per layer; this + checkpoint satisfies both. `global_sf` inverts the largest + `down_proj.input_scale`; the spread of the rest is FC2's `alpha`. + - `input_scale` LINEAR. Swizzled scales have the same size when T % 128 == 0 + and are read as garbage without error. + - The sort from moe_sort with the same `top_k` and tile_tokens_dim == + `tile_size`, over the tokens `input` holds. Nothing checks `top_k` against it. + + Refused by `is_valid`, each silently wrong unguarded: a non-contiguous input, + input_scale, weight, weight_scale or alpha (data_ptr() read as dense); an + alpha shorter than weight.size(0); a weight_scale with fewer bytes than the + weight needs (read past its end: zeros of the same allocation come back as zero + experts, a fresh allocation faults). + + Loud (AssertionError from the runner before the launch): tile_size 256 with a + tile-128 sort and 128 with a tile-256 sort; input_scale not H/16 wide; + global_sf of 2 elements; alpha not fp32; weight_scale typed float8_e4m3fn; + input typed uint8. tile_size 64: ValueError. + + Uncertified: Relu2 and SiTU; a SwiGLU clamp (swiglu_limit_scalar >= 0); + expert_counts (DeepEP count-native metadata); expert parallelism (the kernel + reads neither num_experts nor local_expert_offset, so EP changes only the sort + and the weight slice); other hidden / intermediate sizes and top_k; + scaling_vector_size other than 16; other architectures. + + Hard limits (sm_103) + written rows x I/16 B of c_sf boundary 2^31 B: an illegal memory access + once rows written reach 53687092 at I = 640 (T ~ 5.37M at + top-10; T = 8192 writes at most 212480 rows) + input, input_scale, weight past 2^31 elements / bytes: exact (probed + up to 1536 experts) + permuted index int32, within moe_sort's 2^31-pair limit; no index passes + through fp32 + """ + + def __call__( + self, + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + global_sf: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + scaling_vector_size: int = 16, + activation_type: int = _SWIGLU, + swiglu_limit_scalar: float = SWIGLU_LIMIT_SCALAR_DISABLED, + expert_counts: Optional[torch.Tensor] = None, + expert_capacity: int = 0, + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + return torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell( + input=input, + weight=weight, + input_scale=input_scale, + weight_scale=weight_scale, + alpha=alpha, + tile_idx_to_group_idx=tile_idx_to_group_idx, + tile_idx_to_mn_limit=tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + global_sf=global_sf, + num_experts=num_experts, + top_k=top_k, + num_local_experts=num_local_experts, + local_expert_offset=local_expert_offset, + tile_size=tile_size, + scaling_vector_size=scaling_vector_size, + activation_type=activation_type, + swiglu_limit_scalar=swiglu_limit_scalar, + expert_counts=expert_counts, + expert_capacity=expert_capacity, + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, + ) + + def reference( + self, + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + global_sf: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + scaling_vector_size: int = 16, + activation_type: int = _SWIGLU, + swiglu_limit_scalar: float = SWIGLU_LIMIT_SCALAR_DISABLED, + expert_counts: Optional[torch.Tensor] = None, + expert_capacity: int = 0, + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None, + ) -> Fc1Expected: + """Gather, dequantize, GEMM in float64, SwiGLU, exact NVFP4 rounding. + + Products of e2m1 and e4m3 values are exact, and float64 sums of 2560 of + them are exact to ~1e-16 relative, so this side carries no accumulation + error of its own; the block scale and codes are the exact round-to-nearest- + even of the exact `y * global_sf`. Where a block scale rounds to zero the + codes here are zero: they dequantize to zero whatever they hold, and the + test's compare does not read them. + + Certified surface only: SwiGLU without a clamp, the sort-metadata path + (`expert_counts` None), no SiTU. Anything else raises NotImplementedError. + `num_experts`, `num_local_experts`, `local_expert_offset` and + `expert_capacity` do not enter the computation. + """ + if ( + int(activation_type) != _SWIGLU + or swiglu_limit_scalar >= 0 + or expert_counts is not None + or situ_beta is not None + or situ_linear_beta is not None + ): + raise NotImplementedError( + "reference covers SwiGLU without a clamp on sort metadata only" + ) + vec = scaling_vector_size + dev = input.device + m = permuted_idx_to_expanded_idx.numel() + n2 = weight.shape[1] + inter = n2 // 2 + hidden = input.shape[1] * 2 + n_tiles = int(num_non_exiting_tiles.reshape(-1)[0]) + span = n_tiles * tile_size + + rows = torch.zeros(m, dtype=torch.bool, device=dev) + limit = tile_idx_to_mn_limit[:n_tiles].long().repeat_interleave(tile_size) + rows[:span] = torch.arange(span, device=dev) < limit + real = torch.nonzero(rows).squeeze(1) + token = permuted_idx_to_expanded_idx.long()[real] // top_k + expert = tile_idx_to_group_idx[:n_tiles].long().repeat_interleave(tile_size)[real] + + a = _dequant(input.view(torch.uint8)[token], input_scale.view(torch.uint8)[token], vec) + y = torch.empty(real.numel(), inter, dtype=torch.float64, device=dev) + w_sf = weight_scale.view(torch.uint8).reshape(weight.shape[0], -1) + group = GATE_UP_GROUP + for e in torch.unique(expert).tolist(): + sel = torch.nonzero(expert == e).squeeze(1) + w = _dequant(weight[e].view(torch.uint8), unswizzle(w_sf[e], n2, hidden // vec), vec) + acc = (a[sel] @ w.T) * float(alpha[e]) + acc = acc.view(-1, n2 // (2 * group), 2, group) + up = acc[:, :, 0].reshape(-1, inter) + gate = acc[:, :, 1].reshape(-1, inter) + y[sel] = up * gate * torch.sigmoid(gate) + scaled = y * float(global_sf.reshape(-1)[0]) + + blocks = scaled.view(-1, inter // vec, vec) + pre = blocks.abs().amax(-1) / E2M1_MAX + sf = pre.clamp(max=E4M3_MAX).float().to(torch.float8_e4m3fn) + sf_val = sf.double().unsqueeze(-1) + q = torch.where(sf_val > 0, blocks / sf_val, torch.zeros_like(blocks)) + codes = e2m1_codes(q.reshape(-1, inter)) + c = torch.zeros(m, inter // 2, dtype=torch.uint8, device=dev) + c[real] = codes[:, 0::2] | (codes[:, 1::2] << 4) + sf_rows = torch.zeros(m, inter // vec, dtype=torch.uint8, device=dev) + sf_rows[real] = sf.view(torch.uint8) + return Fc1Expected( + c=c.view(torch.float4_e2m1fn_x2), c_sf=swizzle(sf_rows), rows=rows, scaled=scaled + ) + + def is_valid( + self, + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + global_sf: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + scaling_vector_size: int = 16, + activation_type: int = _SWIGLU, + swiglu_limit_scalar: float = SWIGLU_LIMIT_SCALAR_DISABLED, + expert_counts: Optional[torch.Tensor] = None, + expert_capacity: int = 0, + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None, + ) -> None: + # The runner hands data_ptr() to the kernel with dense row-major layouts + # assumed; it checks dims and dtypes, never strides. + for name, t in ( + ("input", input), + ("input_scale", input_scale), + ("weight", weight), + ("weight_scale", weight_scale), + ("alpha", alpha), + ): + assert t.is_contiguous(), ( + f"{name} must be contiguous; the kernel reads it from data_ptr() as a " + "dense row-major buffer and silently reads the wrong elements" + ) + # One alpha per expert of `weight`; the runner checks only alpha.dim() == 1. + assert alpha.numel() == weight.size(0), ( + f"alpha has {alpha.numel()} elements for {weight.size(0)} experts; the " + "kernel reads alpha[expert] past the end without raising" + ) + # The runner checks weight_scale.dim() == 3 only, never its size. + expected = weight.numel() * 2 // scaling_vector_size + assert weight_scale.numel() == expected, ( + f"weight_scale holds {weight_scale.numel()} bytes, the weight needs " + f"{expected}; the kernel reads past its end -- silently within the " + "allocation, with an illegal memory access beyond it" + ) + + +# ── native-torch NVFP4 helpers for reference() and the test's compare ──────── + + +def unswizzle(buf: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + """128x4-swizzled scale bytes -> `[rows, cols]` (rows % 128 == 0, cols % 4 == 0). + + Byte `(r, c)` sits at `(r // 128) * 128 * cols + (c // 4) * 512 + (r % 32) * 16 + + ((r % 128) // 32) * 4 + c % 4`. + """ + return buf.reshape(rows // 128, cols // 4, 32, 4, 4).permute(0, 3, 2, 1, 4).reshape(rows, cols) + + +def swizzle(sf: torch.Tensor) -> torch.Tensor: + """`[rows, cols]` scale bytes -> the flat 128x4-swizzled buffer; inverse of `unswizzle`.""" + rows, cols = sf.shape + return sf.reshape(rows // 128, 4, 32, cols // 4, 4).permute(0, 3, 2, 1, 4).reshape(-1) + + +def e4m3_value(byte: torch.Tensor) -> torch.Tensor: + return byte.to(torch.uint8).view(torch.float8_e4m3fn).double() + + +def _dequant(packed: torch.Tensor, sf: torch.Tensor, vec: int) -> torch.Tensor: + """`[R, K/2]` packed e2m1 (element 2j in the low nibble) and `[R, K/vec]` e4m3 + bytes -> float64 `[R, K]`. Exact.""" + lut = torch.tensor(E2M1, dtype=torch.float64, device=packed.device) + lo, hi = (packed & 0xF).long(), (packed >> 4).long() + codes = torch.stack([lo, hi], dim=-1).reshape(packed.shape[0], -1) + vals = lut[codes & 7] * torch.where(codes >= 8, -1.0, 1.0).double() + return vals * e4m3_value(sf).repeat_interleave(vec, dim=1) + + +def e2m1_codes(v: torch.Tensor) -> torch.Tensor: + """float -> 4-bit e2m1 codes (bit 3 = sign), round-to-nearest-even, saturating at 6.""" + mid = torch.tensor(_E2M1_MIDPOINTS, device=v.device, dtype=v.dtype) + tie_up = torch.tensor(_TIE_UP, device=v.device) + a = v.abs().unsqueeze(-1) + code = ((a > mid).sum(-1) + ((a == mid) & tie_up).any(-1)).to(torch.uint8) + negative = (v < 0) | ((v == 0) & torch.signbit(v)) + return code | (negative.to(torch.uint8) << 3) + + +cute_dsl_nvfp4_gather_grouped_gemm_act_fusion = CuteDslNvfp4GatherGroupedGemmActFusion() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_grouped_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_grouped_gemm.py new file mode 100644 index 000000000000..81216c78ecf0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_grouped_gemm.py @@ -0,0 +1,457 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CuteDSL NVFP4 FC2 of the routed experts: the plain grouped GEMM over the permuted +FC1 output, one bf16 row per permuted row, in one trtllm call.""" + +from __future__ import annotations + +from typing import NamedTuple + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: e2m1 magnitudes by 3-bit code. +_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) + + +def _ckpt( + tokens: int, + tile_size: int, + mma_n: int = 128, + cluster_n: int = 1, + dist: str = "router", + gsf: str = "calibrated", + zero_tokens: int = 0, +) -> dict: + """The checkpoint's routed experts: 512 experts, top-10, hidden 2560, I 640, all + local (tp1), bf16 out, 16-element e4m3 block scales. + + `mma_n` / `cluster_n` select the autotuner tactic: MMA tile `(tile_size, + mma_n)`, cluster `(tile_size // 128, cluster_n)`; (128, 1) is also the + fallback when the tuner has no entry. `dist` names how the test draws the + expert ids, as in the moe_sort entry's test: `router` = each row's top-10 of + uniform scores; `hot` = every row on the same 10 experts; `tile_edges` = three + experts hold exactly tile-1, tile and tile+1 rows; `round_robin` = row i + takes experts (10 i + j) mod 512. `gsf` is the FC1 `global_sf` that quantized + this op's input, against the SwiGLU output's magnitude: `calibrated` = the + output maximum the test calibrates to maps to the e4m3 maximum, as the + checkpoint's per-layer `1 / down_proj.input_scale` does; `saturating` = 25x + above that, so most input block scales clip at 448 and the codes at 6; + `subnormal` = 2^-12 of it, so input block scales fall into e4m3 subnormals + and zero. `zero_tokens` makes that many leading tokens' rows all zero, as + FC1 returns them for an all-zero token. + """ + return dict( + tokens=tokens, + tile_size=tile_size, + mma_n=mma_n, + cluster_n=cluster_n, + dist=dist, + gsf=gsf, + zero_tokens=zero_tokens, + num_experts=512, + top_k=10, + hidden=2560, + intermediate=640, + ) + + +class Fc2Expected(NamedTuple): + """What `reference` returns: the op's output, plus what the test's compare reads. + + `out` is laid out exactly as the op returns it; rows the op does not write + hold zeros. `rows` marks the rows the op writes -- always the prefix of + `num_non_exiting_tiles * tile_size` rows. `exact` is `alpha[e] * sum_k a_k + w_k` on those rows, before any rounding; `magnitude` is `|alpha[e]| * sum_k + |a_k w_k|`, the scale every fp32 rounding of the accumulation is relative to. + """ + + out: torch.Tensor + rows: torch.Tensor + exact: torch.Tensor + magnitude: torch.Tensor + + +class CuteDslNvfp4GroupedGemm(OpWrapper): + """`torch.ops.trtllm.cute_dsl_nvfp4_grouped_gemm_blackwell`: FC2 of the routed + experts on the CuteDSL MoE path, one kernel launch. + + For every row `p` of the first `num_non_exiting_tiles` tiles (tile `i = p // + tile_size`): + + e = tile_idx_to_group_idx[i] the local expert + acc[p] = deq(input[p]) @ deq(weight[e]).T [H], fp32 + out[p] = bf16(alpha[e] * acc[p]) fp32 product + + where `deq` is e2m1 times its e4m3 block scale. Every row of a live tile is + computed, padding rows included; no row is gathered, weighted or summed. + + Arguments, as the CuteDSL MoE passes them (`run_moe_nvfp4_impl`, + `_torch/moe/fused_moe/fused_moe_cute_dsl.py`, its non-fused-finalize branch): + + - `input` `[M, I/2]` float4_e2m1fn_x2 and `input_scale` `[M * I/16]` uint8: + FC1's output exactly as the CuteDSL FC1 + (`cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell`) returns it -- + one row per PERMUTED row, the block scales 128x4-swizzled over `[M, + I/16]`. `M = permuted_idx_to_expanded_idx.numel()` of the sort. + - `weight` `[E, H, I/2]` float4_e2m1fn_x2: `down_proj` per expert in the + checkpoint's row order. No 32-row shuffle (that is trtllm-gen's layout). + - `weight_scale` `[E, H, I/16]` uint8: its e4m3 block scales, 128x4-swizzled + per expert (`block_scale_interleave`). + - `alpha` `[E]` fp32: per expert, `max(down_proj.input_scale) * weight_scale_2[e]` + -- the inverse of FC1's `global_sf` times the expert's weight scale. + - `tile_idx_to_group_idx` `[M / tile_size]` int32 and `num_non_exiting_tiles` + `[1]` int32: moe_sort's `tile_idx_to_expert_idx` and + `num_non_exiting_tiles` from the sort FC1 ran on, with tile_tokens_dim == + `tile_size`. + - `num_experts`, `top_k`, `num_local_experts`, `local_expert_offset`: read only + by the autotuner's shape buckets and cache key; the kernel takes the + expert count from `weight.size(0)` and the expert from the tile. + - `output_dtype`: bf16, the only one the runner builds. + + Returns `[M, H]` bfloat16, one row per permuted row. The pair `(t, k)` is row + `expanded_idx_to_permuted_idx[t, k]`; rows outside the live tiles are not + written (see `note`). + + Fusion boundary: the grouped GEMM, the per-expert alpha and the bf16 + conversion happen inside. The rows arrive already permuted (FC1 gathered + them); the routing-weighted top-k combine is the caller's + (`moe_unpermute_inplace`). Routing weights are not read here. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = tuple( + cell + for tile in (128, 256) + for cell in ( + Cell( + why=f"tile {tile}: decode, 1 request without MTP", + spec=_ckpt(1, tile), + ), + Cell( + why=f"tile {tile}: 1 request at MTP verify width 2", + spec=_ckpt(2, tile), + ), + Cell( + why=f"tile {tile}: 1 request at MTP verify width 4", + spec=_ckpt(4, tile), + ), + Cell( + why=f"tile {tile}: decode, 8 requests", + spec=_ckpt(8, tile), + ), + Cell( + why=f"tile {tile}: decode at the 32-request ceiling", + spec=_ckpt(32, tile), + ), + Cell( + why=f"tile {tile}: 32 requests at MTP verify width 3", + spec=_ckpt(96, tile), + ), + Cell( + why=f"tile {tile}: 32 requests at MTP verify width 4; largest captured T", + spec=_ckpt(128, tile), + ), + Cell( + why=f"tile {tile}: 8192 rows, max_num_tokens prefill", + spec=_ckpt(8192, tile), + ), + Cell( + why=f"tile {tile}: 8192 rows all on 10 experts; {8192 // tile} full tiles each", + spec=_ckpt(8192, tile, dist="hot"), + ), + Cell( + why=f"tile {tile}: expert rows tile-1, tile, tile+1; the tile edge", + spec=_ckpt(300, tile, dist="tile_edges"), + ), + Cell( + why=f"tile {tile}: 510 experts with one row each; every tile mostly padding", + spec=_ckpt(51, tile, dist="round_robin"), + ), + Cell( + why=f"tile {tile}: MMA N 256 tactic, decode", + spec=_ckpt(1, tile, mma_n=256), + ), + Cell( + why=f"tile {tile}: MMA N 256 tactic, 128 rows", + spec=_ckpt(128, tile, mma_n=256), + ), + Cell( + why=f"tile {tile}: MMA N 256 tactic, 8192-row prefill", + spec=_ckpt(8192, tile, mma_n=256), + ), + Cell( + why=f"tile {tile}: cluster N 2 tactic, 128 rows", + spec=_ckpt(128, tile, cluster_n=2), + ), + Cell( + why=f"tile {tile}: MMA N 256 and cluster N 2 tactic, 8192-row prefill", + spec=_ckpt(8192, tile, mma_n=256, cluster_n=2), + ), + Cell( + why=f"tile {tile}: FC1 global_sf 25x calibrated; input scales clip at 448", + spec=_ckpt(2048, tile, gsf="saturating"), + ), + Cell( + why=f"tile {tile}: tiny FC1 global_sf; subnormal and zero input scales", + spec=_ckpt(2048, tile, gsf="subnormal"), + ), + Cell( + why=f"tile {tile}: 7 all-zero token rows; zero input rows", + spec=_ckpt(64, tile, zero_tokens=7), + ), + ) + ) + + note = """ + Path certified: the CuTe DSL kernel `Sm100BlockScaledContiguousGroupedGemmKernel` + (`_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_grouped_gemm.py`), + built by `Sm100BlockScaledContiguousGroupedGemmRunner` + (`_torch/custom_ops/cute_dsl_custom_ops.py`), which JIT-compiles it with + `cute.compile` on the first call per (tile, tactic) and keeps it for the + process. The op is registered only when the CuTe DSL imports + (`IS_CUTLASS_DSL_AVAILABLE`), and the runner accepts SM 100 and 103 only. + + Metadata consumed: none beyond the arguments. Stateless: allocates its output + with torch.empty on every call, on the current stream, with no host + synchronization; `num_non_exiting_tiles` is read on the device, so it + captures into a CUDA graph together with moe_sort and FC1. + + Call sites: the engine's CuteDSL MoE (`run_moe_nvfp4_impl`, when the fused + finalize is off) and the target's routed-expert step, both followed by + `moe_unpermute_inplace`; eager for prefill and mixed batches (T up to + max_num_tokens 8192) and captured for decode (T = requests x MTP verify + width, up to 32 x 4 = 128). Tile 128 is the target's; 256 is certified as + well. + + Numerics: each written element equals the reference's `out` bit for bit on + every cell, on the real FC1 output at 8192 tokens, and on the checkpoint's own + down_proj experts. Operands whose block scales span a wider range can round a + partial sum; the test's compare admits that at bf16 ties only. + + Tactics. The op asks the AutoTuner for an MMA tile of (tile_size, 128) or + (tile_size, 256) and a cluster of (tile_size // 128, 1) or (tile_size // + 128, 2), and with no tuning entry falls back to ((tile_size, 128), + (tile_size // 128, 1)); serving warmup tunes per row bucket and may pick any + of the four. Every pair's row is bitwise identical across tile 128 / 256 x + the four tactics x the fallback, so cells driven at one tactic certify the + others. + + Determinism. A (token, expert) pair's output row depends on that pair's input + row and its expert alone: it is bitwise unchanged by moving it to another row + of its expert's run (moe_sort orders them by atomics from T = 5), by repeating + the call, and by cutting a batch to a prefix. This op does not make the MoE + run-to-run nondeterministic; the fused-finalize variant + (`cute_dsl_nvfp4_grouped_gemm_finalize_inplace_blackwell`) reduces across a + token's experts in launch order and does. + + What is read and written: every row of the first num_non_exiting_tiles tiles + is written; the pair rows hold their results and the padding rows of a live + tile are computed from whatever the input holds there -- zero for FC1's + output, which zeroes them. Rows at and past num_non_exiting_tiles * tile_size + are never written: eager they hold whatever torch.empty returned, under a + graph whatever the pool last held. Input rows past that point reach no written + row. The combine must read only expanded_idx_to_permuted_idx's rows. + + Preconditions the caller owns: + - `weight` / `weight_scale` in the layout the class docstring states, and + `input_scale` swizzled. Each wrong preparation ran without error, with + finite output, and put most of the pair rows' elements outside the test's + gate: trtllm-gen's 32-row shuffle, weight scales left linear, input scales + linear, alpha shifted by one expert, one alpha for the whole layer. + - `alpha` per expert, `max(down_proj.input_scale) * weight_scale_2[e]`, with + the maximum the value FC1's `global_sf` inverts. + - `tile_idx_to_group_idx` / `num_non_exiting_tiles` from the same sort as + FC1's input, tile_tokens_dim == `tile_size`, and every expert index below + `weight.size(0)`: nothing checks the values (a tile naming expert 512 of + 512 ran without error and came back all zeros). + - Every tensor on the device. Nothing checks it: a host-resident `alpha` or + tile table was read correctly on a coherent-memory Grace host, unmeasured + where the GPU cannot dereference a host pointer. + + Refused by `is_valid`, each silently wrong unguarded: a non-contiguous input, + input_scale, weight, weight_scale, alpha or tile_idx_to_group_idx (data_ptr() + read as dense; finite, wrong rows). + + Loud (AssertionError from the runner before the launch): a tile_size other + than the sort's, either way round; output_dtype float16; alpha bf16 or short + of weight.size(0); weight_scale with fewer experts or typed float8_e4m3fn; + input typed uint8; input_scale 2-D or short; scaling_vector_size 32; + num_non_exiting_tiles of 2 elements; tile_idx_to_group_idx int64 or one tile + short. tile_size 64: ValueError. + + Uncertified: output dtypes other than bf16 (the runner builds none); expert + parallelism (the kernel reads neither local_expert_offset nor num_experts, so + EP changes only the sort and the weight slice); other hidden / intermediate + sizes; scaling_vector_size other than 16; other architectures. + + Hard limits (sm_103) + input, input_scale, output, weight, weight_scale past 2^31 elements / + bytes: exact (probed up to M = 53687296 rows and 20973 + experts) + tile index int32 tables, within moe_sort's 2^31-pair limit; no index + passes through fp32 + """ + + def __call__( + self, + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_dtype: torch.dtype, + scaling_vector_size: int = 16, + ) -> torch.Tensor: + return torch.ops.trtllm.cute_dsl_nvfp4_grouped_gemm_blackwell( + input=input, + weight=weight, + input_scale=input_scale, + weight_scale=weight_scale, + alpha=alpha, + tile_idx_to_group_idx=tile_idx_to_group_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + num_experts=num_experts, + top_k=top_k, + num_local_experts=num_local_experts, + local_expert_offset=local_expert_offset, + tile_size=tile_size, + output_dtype=output_dtype, + scaling_vector_size=scaling_vector_size, + ) + + def reference( + self, + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_dtype: torch.dtype, + scaling_vector_size: int = 16, + ) -> Fc2Expected: + """Dequantize, GEMM in float64, scale by alpha in fp32, round once to bf16. + + Products of e2m1 and e4m3 values are exact, and float64 sums of 640 of + them are exact to ~1e-16 relative, so `exact` carries no accumulation + error of its own; `out` is that sum rounded to fp32, multiplied by the + fp32 alpha, rounded to bf16 -- what an fp32-accumulating GEMM with this + epilogue returns whenever its fp32 partial sums are exact. Rows whose + input dequantizes to all zeros are zero without a GEMM. + + bf16 output only (anything else raises NotImplementedError). + `num_experts`, `top_k`, `num_local_experts` and `local_expert_offset` do + not enter the computation. + """ + if output_dtype != torch.bfloat16: + raise NotImplementedError("reference covers bf16 output only") + vec = scaling_vector_size + dev = input.device + m = input.shape[0] + hidden = weight.shape[1] + k_blocks = input.shape[1] * 2 // vec + n_tiles = int(num_non_exiting_tiles.reshape(-1)[0]) + span = n_tiles * tile_size + + rows = torch.zeros(m, dtype=torch.bool, device=dev) + rows[:span] = True + a = _dequant( + input.view(torch.uint8)[:span], + unswizzle(input_scale.view(torch.uint8), m, k_blocks)[:span], + vec, + ) + live = a.abs().amax(dim=1) > 0 + expert = tile_idx_to_group_idx[:n_tiles].long().repeat_interleave(tile_size) + acc = torch.zeros(span, hidden, dtype=torch.float64, device=dev) + mag = torch.zeros(span, hidden, dtype=torch.float64, device=dev) + w_sf = weight_scale.view(torch.uint8).reshape(weight.shape[0], -1) + for e in torch.unique(expert[live]).tolist(): + sel = torch.nonzero(live & (expert == e)).squeeze(1) + w = _dequant(weight[e].view(torch.uint8), unswizzle(w_sf[e], hidden, k_blocks), vec) + acc[sel] = a[sel] @ w.T + mag[sel] = a[sel].abs() @ w.abs().T + scale = alpha.float()[expert].unsqueeze(1) + out = torch.zeros(m, hidden, dtype=torch.bfloat16, device=dev) + out[:span] = (acc.float() * scale).to(torch.bfloat16) + return Fc2Expected( + out=out, + rows=rows, + exact=acc * scale.double(), + magnitude=mag * scale.double().abs(), + ) + + def is_valid( + self, + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_dtype: torch.dtype, + scaling_vector_size: int = 16, + ) -> None: + # The runner hands data_ptr() to the kernel with dense row-major layouts + # assumed; it checks dims, sizes and dtypes, never strides. + for name, t in ( + ("input", input), + ("input_scale", input_scale), + ("weight", weight), + ("weight_scale", weight_scale), + ("alpha", alpha), + ("tile_idx_to_group_idx", tile_idx_to_group_idx), + ): + assert t.is_contiguous(), ( + f"{name} must be contiguous; the kernel reads it from data_ptr() as a " + "dense row-major buffer and silently reads the wrong elements" + ) + + +# ── native-torch NVFP4 helpers for reference() and the test's compare ──────── + + +def unswizzle(buf: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + """128x4-swizzled scale bytes -> `[rows, cols]` (rows % 128 == 0, cols % 4 == 0). + + Byte `(r, c)` sits at `(r // 128) * 128 * cols + (c // 4) * 512 + (r % 32) * 16 + + ((r % 128) // 32) * 4 + c % 4`. + """ + return buf.reshape(rows // 128, cols // 4, 32, 4, 4).permute(0, 3, 2, 1, 4).reshape(rows, cols) + + +def _dequant(packed: torch.Tensor, sf: torch.Tensor, vec: int) -> torch.Tensor: + """`[R, K/2]` packed e2m1 (element 2j in the low nibble) and `[R, K/vec]` e4m3 + bytes -> float64 `[R, K]`. Exact.""" + nibble = torch.tensor(_E2M1, dtype=torch.float64, device=packed.device) + nibble = torch.cat([nibble, -nibble]) + byte = torch.arange(256, device=packed.device) + pair = torch.stack([nibble[byte & 0xF], nibble[byte >> 4]], dim=1) + vals = pair[packed.long()].reshape(packed.shape[0], -1) + scales = sf.to(torch.uint8).view(torch.float8_e4m3fn).double() + return vals * scales.repeat_interleave(vec, dim=1) + + +cute_dsl_nvfp4_grouped_gemm = CuteDslNvfp4GroupedGemm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_sort.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_sort.py new file mode 100644 index 000000000000..8aa56756273e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_sort.py @@ -0,0 +1,308 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MoE sort: bucket the routed (token, expert) pairs into per-expert row tiles +for the CuteDSL grouped GEMMs, in one trtllm call.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +def _ckpt(tokens: int, tile_tokens_dim: int, dist: str = "router") -> dict: + """The checkpoint's trunk MoE: 512 experts, top-10, all local (tp1), fp32 scales. + + `dist` names how the test draws the expert ids: `router` = each row's top-10 of + uniform scores over the 512 experts; `hot` = every row routes to the same 10 + experts; `tile_edges` = three experts hold exactly tile-1, tile and tile+1 rows, + the rest `router`; `round_robin` = row i takes experts (10 i + j) mod 512; + `strided` = row i takes (h_i + 51 j) mod 512 with h_i uniform. All of them put 10 + distinct experts on every row, as the router does. + """ + return dict( + tokens=tokens, + dist=dist, + num_experts=512, + top_k=10, + local_expert_offset=0, + local_num_experts=512, + tile_tokens_dim=tile_tokens_dim, + scales_dtype=torch.float32, + ) + + +class MoeSort(OpWrapper): + """`torch.ops.trtllm.moe_sort`: routed (token, slot) pairs -> per-expert row tiles. + + Input: `token_selected_experts` `[T, top_k]` int32 global expert ids (the + router's output), `token_final_scales` `[T, top_k]` (its weights). Call it + with the op's own keyword names, as the CuteDSL MoE does. + + Numbering: an *expanded* index is `t * top_k + k` (row-major over + `token_selected_experts`); a *permuted* index is a row of the grouped GEMM's + M dimension. Each local expert `e = id - local_expert_offset` owns a + contiguous run of `count_e` permuted rows starting on a tile boundary, + experts in increasing order, the run padded up to whole tiles of + `tile_tokens_dim` rows. Returns a list of six int32 tensors, in this order: + + - `tile_idx_to_expert_idx` `[max_tiles]`: the local expert of tile `i`; + - `tile_idx_to_mn_limit` `[max_tiles]`: one past the last real row of tile + `i`, as an absolute permuted index: `min((i + 1) * tile, start_e + count_e)`; + - `expanded_idx_to_permuted_idx` `[T, top_k]`: the permuted row of each + pair, `-1` for a pair whose expert is not local; + - `permuted_idx_to_expanded_idx` `[max_tiles * tile]`: the inverse on real rows; + - `total_num_padded_tokens` `[1]` = `num_non_exiting_tiles * tile`; + - `num_non_exiting_tiles` `[1]`: tiles actually used. + + `max_tiles = m + (T * top_k - m) // tile` with `m = min(local_num_experts, + T * top_k)` is computed on the host from the shapes and scalars alone, so + every output shape is fixed by `T`, which is what lets a CUDA graph capture + it. Only the first `num_non_exiting_tiles` entries of the tile arrays and + the rows below each tile's `mn_limit` are written; everything else is left + as allocated (see `note`). + + Fusion boundary: sorting only. Routing (the top-k) is the caller's, and so + are the permutation of the activations (the FC1 gather does it) and the + routing-weighted unpermute; this op does not read `token_final_scales` into + any output. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = tuple( + cell + for tile in (128, 256) + for cell in ( + Cell( + why=f"tile {tile}: decode, 1 request without MTP; block kernel", + spec=_ckpt(1, tile), + ), + Cell( + why=f"tile {tile}: 4 rows, 1 request at MTP verify width 4; last block-kernel T", + spec=_ckpt(4, tile), + ), + Cell( + why=f"tile {tile}: 5 rows; first single-cluster-kernel T", + spec=_ckpt(5, tile), + ), + Cell( + why=f"tile {tile}: 32 rows, decode at the 32-request ceiling", + spec=_ckpt(32, tile), + ), + Cell( + why=f"tile {tile}: 128 rows, 32 requests at MTP verify width 4; largest captured T", + spec=_ckpt(128, tile), + ), + Cell( + why=f"tile {tile}: 8192 rows, max_num_tokens prefill; last single-cluster T", + spec=_ckpt(8192, tile), + ), + Cell( + why=f"tile {tile}: 8192 rows all to the same 10 experts; full multi-tile runs", + spec=_ckpt(8192, tile, "hot"), + ), + Cell( + why=f"tile {tile}: expert counts tile-1, tile, tile+1; the mn_limit clamp", + spec=_ckpt(300, tile, "tile_edges"), + ), + Cell( + why=f"tile {tile}: 510 experts with one row each; n_tiles == max_tiles", + spec=_ckpt(51, tile, "round_robin"), + ), + Cell( + why=f"tile {tile}: 8193 rows, past the envelope; first cooperative-kernel T", + spec=_ckpt(8193, tile), + ), + Cell( + why=f"tile {tile}: 471859 rows; last cooperative-kernel T on 152 SMs", + spec=_ckpt(471859, tile, "strided"), + ), + Cell( + why=f"tile {tile}: 471860 rows; first histogram+offsets T on 152 SMs", + spec=_ckpt(471860, tile, "strided"), + ), + ) + ) + + note = """ + Metadata consumed: none. Stateless; allocates its six outputs plus scratch on + every call and runs on the current stream, with no host synchronization, so + it captures into a CUDA graph. + + Call it as the engine's CuteDSL MoE does (`run_moe_nvfp4_impl`, + `_torch/moe/fused_moe/fused_moe_cute_dsl.py`): `num_experts` = the slot + count, `local_expert_offset` / `local_num_experts` = this rank's slots, and + `tile_tokens_dim` = the grouped GEMMs' tile, 128 or 256 on sm_103. The + downstream FC1 gather GEMM, output memset and FC2 take these six tensors as + they are. + + Kernel path, by `T`: + T 1..4 routingCustom::routingIndicesBlockKernel + T 5..8192 routingCustom::routingIndicesClusterKernel + T 8193..471859 routingInitExpertCounts + routingIndicesCoopKernel + T >= 471860 routingInitExpertCounts + routingIndicesHistogramKernel + + routingIndicesOffsetsKernel + The cooperative ceiling is (SMs - 8) * 512 * 64 / top_k rows with 512 + experts: 471859 on a 152-SM GB300, a different T on another SM count. + + Order within an expert's run is unspecified. The block kernel (T <= 4) lays + each expert's rows out in token order and is bitwise repeatable; every other + path assigns them by atomics, so `expanded_idx_to_permuted_idx` differs + between identical calls. A test therefore compares each expert's rows as a + set and the two mappings as exact inverses; the tile arrays and counts are + deterministic. + + Not written: the tile arrays past `num_non_exiting_tiles`, and the rows of + `permuted_idx_to_expanded_idx` from each tile's `mn_limit` to the tile's + end. Eager, those hold whatever the allocator handed back; under a CUDA + graph they hold the previous replay's indices, which look valid. Every + consumer must bound its reads by `num_non_exiting_tiles` and + `tile_idx_to_mn_limit`. + + `token_final_scales`: shape-checked and never written; its values never + reach an output. Its pointer must be non-null, which is why T = 0 raises. + + Data hazards `is_valid` cannot see (values): a row naming the same expert + twice keeps both pairs on the cluster kernel but loses one on the block + kernel (that pair gets -1). An id outside [local_expert_offset, + local_expert_offset + local_num_experts) gets -1, except that the cluster + kernel reads ids through int16, so an id of 65536 + e sorts as expert e. The + router produces neither. + + Loud: int64 ids, 1-D ids, a `[T, top_k]` size mismatch on either input, + CPU tensors, top_k above 32, num_experts above 2048, T = 0. + + Silent, refused by `is_valid`: a non-contiguous `token_selected_experts` + (read from `data_ptr()` as dense `[T, top_k]`); 2^31 or more elements (the + host counts pairs in int32: it raises a negative-dimension error from 2^31 + and silently sorts only the wrapped remainder from 2^32); `tile_tokens_dim` + <= 0 (the host divides by it). + + Certified only at the checkpoint's values: num_experts 512, top_k 10, all + experts local (offset 0), tile 128 and 256, fp32 scales. Expert parallelism, + other tiles and other expert counts are uncertified. + + Hard limits (sm_103) + T * top_k boundary 2^31 pairs (int32 product) + max_tiles * tile boundary 2^31 padded rows (raises past it) + expert id int16 in the cluster kernel; ids < 2048 never reach it + """ + + def __call__( + self, + token_selected_experts: torch.Tensor, + token_final_scales: torch.Tensor, + num_experts: int, + top_k: int, + local_expert_offset: int, + local_num_experts: int, + tile_tokens_dim: int, + ) -> list[torch.Tensor]: + return torch.ops.trtllm.moe_sort( + token_selected_experts=token_selected_experts, + token_final_scales=token_final_scales, + num_experts=num_experts, + top_k=top_k, + local_expert_offset=local_expert_offset, + local_num_experts=local_num_experts, + tile_tokens_dim=tile_tokens_dim, + ) + + def reference( + self, + token_selected_experts: torch.Tensor, + token_final_scales: torch.Tensor, + num_experts: int, + top_k: int, + local_expert_offset: int, + local_num_experts: int, + tile_tokens_dim: int, + ) -> list[torch.Tensor]: + """Counting sort by local expert, stable in the expanded index. + + Within an expert, rows are in increasing expanded index; the test's compare + accepts any order there. What the op leaves unwritten is `-1` here. + `token_final_scales`, `num_experts` and `top_k` are not read: the scales + never reach an output, and the ids' shape carries the rest. + """ + ids = token_selected_experts + dev = ids.device + tile = tile_tokens_dim + num_tokens, width = ids.shape + num_expanded = num_tokens * width + local = ids.reshape(-1).long() - local_expert_offset + is_local = (local >= 0) & (local < local_num_experts) + counts = torch.bincount(local[is_local], minlength=local_num_experts) + tiles_per_expert = (counts + tile - 1) // tile + num_tiles = int(tiles_per_expert.sum()) + first = min(local_num_experts, num_expanded) + max_tiles = first + (num_expanded - first) // tile + + run_start = (torch.cumsum(tiles_per_expert, 0) - tiles_per_expert) * tile + tile_expert = torch.repeat_interleave( + torch.arange(local_num_experts, device=dev), tiles_per_expert + ) + tile_idx = torch.arange(num_tiles, device=dev) + mn_limit = torch.minimum( + (tile_idx + 1) * tile, run_start[tile_expert] + counts[tile_expert] + ) + + expanded = torch.arange(num_expanded, device=dev)[is_local] + expert = local[is_local] + order = torch.sort(expert, stable=True).indices + expanded, expert = expanded[order], expert[order] + rank = torch.arange(expert.numel(), device=dev) - (torch.cumsum(counts, 0) - counts)[expert] + permuted = run_start[expert] + rank + + def unwritten(n: int) -> torch.Tensor: + return torch.full((n,), -1, dtype=torch.long, device=dev) + + tile_idx_to_expert_idx = unwritten(max_tiles) + tile_idx_to_expert_idx[:num_tiles] = tile_expert + tile_idx_to_mn_limit = unwritten(max_tiles) + tile_idx_to_mn_limit[:num_tiles] = mn_limit + permuted_idx_to_expanded_idx = unwritten(max_tiles * tile) + permuted_idx_to_expanded_idx[permuted] = expanded + expanded_idx_to_permuted_idx = unwritten(num_expanded) + expanded_idx_to_permuted_idx[expanded] = permuted + outs = ( + tile_idx_to_expert_idx, + tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx.view(num_tokens, width), + permuted_idx_to_expanded_idx, + torch.tensor([num_tiles * tile], device=dev), + torch.tensor([num_tiles], device=dev), + ) + return [t.to(torch.int32) for t in outs] + + def is_valid( + self, + token_selected_experts: torch.Tensor, + token_final_scales: torch.Tensor, + num_experts: int, + top_k: int, + local_expert_offset: int, + local_num_experts: int, + tile_tokens_dim: int, + ) -> None: + # The host computes T * top_k in int32: from 2^31 the op raises a + # negative-dimension error, from 2^32 it sorts only the wrapped remainder. + numel = token_selected_experts.numel() + assert numel < 2**31, ( + f"token_selected_experts has {numel} >= 2^31 elements; moe_sort counts " + "pairs in int32 and silently sorts only the wrapped remainder past 2^32" + ) + # The kernels read the ids from data_ptr() as a dense [T, top_k] buffer. + assert token_selected_experts.is_contiguous(), ( + "token_selected_experts must be contiguous; a strided view is read as a " + "dense [T, top_k] buffer and silently sorts the wrong ids" + ) + assert tile_tokens_dim > 0, ( + f"tile_tokens_dim must be positive, got {tile_tokens_dim}; 0 returns a " + "negative num_non_exiting_tiles without raising" + ) + + +moe_sort = MoeSort() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_unpermute_inplace.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_unpermute_inplace.py new file mode 100644 index 000000000000..24ceefb36d4f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_unpermute_inplace.py @@ -0,0 +1,285 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MoE unpermute: each token's routing-weighted sum of its top-k expert rows, taken +from the grouped GEMM's permuted output in a fixed slot order, in one trtllm call. +""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +def _ckpt( + tokens: int, + tile_tokens_dim: int = 128, + dist: str = "router", + values: str = "normal", + non_local_from: int | None = None, +) -> dict: + """The checkpoint's routed-expert combine: top-10 of 512 experts, hidden 2560, + bf16 FC2 rows, fp32 routing weights. + + The test lays `permuted_input` out exactly as the CuteDSL FC2 receives it: it + sorts the ids with `moe_sort` at `tile_tokens_dim` and writes each (token, + slot) row at the permuted index the sort gave it. `dist` names how the ids are + drawn, as in the moe_sort entry's test: `router` = each row's top-10 of uniform + scores; `hot` = every row on the same 10 experts; `round_robin` = row i takes + experts (10 i + j) mod 512. `values` is how the FC2 rows are drawn: `normal` = + standard normal; `wide` = standard normal times 2^u, u uniform in [-30, 30] + per element, so one sum adds terms up to 2^60 apart in magnitude. + `non_local_from` marks every slot whose expert is at or above it `-1`, as + moe_sort does for a pair on another rank's expert, and also token 0's ten + slots. The scales are a softmax over each row's ten slots, as the + renormalizing router emits them. + """ + return dict( + tokens=tokens, + tile_tokens_dim=tile_tokens_dim, + dist=dist, + values=values, + non_local_from=non_local_from, + num_experts=512, + top_k=10, + hidden_size=2560, + dtype=torch.bfloat16, + scales_dtype=torch.float32, + ) + + +def _fma_f32(acc: torch.Tensor, prod: torch.Tensor) -> torch.Tensor: + """`acc + prod` rounded once to fp32, as a fused multiply-add rounds it. + + `acc` is fp32; `prod` is the fp64 product of a bf16 and an fp32 value, which is + exact (at most 8 + 24 significant bits). The fp64 sum is rounded to odd -- its + exact rounding error from TwoSum, then the last bit forced to 1 if the sum was + inexact -- so that the conversion to fp32 is one correct rounding of the exact + sum rather than two roundings in a row (53 >= 24 + 2 bits). + """ + a = acc.double() + s = a + prod + bb = s - a + err = (a - (s - bb)) + (prod - bb) + inexact = (err != 0) & torch.isfinite(s) + # A sum rounded away from zero truncates to its neighbour one step toward zero. + toward_zero = inexact & ((err < 0) != (s < 0)) + bits = s.view(torch.int64) + bits = torch.where(inexact, (bits - toward_zero.long()) | 1, bits) + return bits.view(torch.float64).float() + + +class MoeUnpermuteInplace(OpWrapper): + """`torch.ops.trtllm.moe_unpermute_inplace`: the routing-weighted top-k combine. + + For each token `t`, writes row `t` of `output`: + + output[t] = bf16( sum over k = 0 .. top_k - 1, in that order, of + permuted_input[expanded_idx_to_permuted_idx[t, k]] + * topk_scales[t, k] ) + + accumulated in fp32 with one rounding per slot (each product fused into the + running sum) and rounded once to the output dtype. A slot whose permuted index + is negative is skipped; a token with every slot skipped gets a zero row. + + Inputs, as the CuteDSL MoE hands them over: `permuted_input` `[M, H]` bf16 is + the FC2 grouped GEMM's output, one row per (token, expert) pair in permuted + order, `M >= T * top_k` (the sort's padded row count); `output` `[T, H]` bf16 + is written in full; `expanded_idx_to_permuted_idx` `[T, top_k]` int32 is the + third output of `moe_sort`; `topk_scales` `[T, top_k]` fp32 are the routing + weights. Call it with these keyword names, the op's own. Returns `output`. + + Fusion boundary: the gather, the weighting and the sum. The FC2 GEMM and the + sort are the caller's, and so is anything added to the routed output (the + shared expert). The op is deterministic: one thread block per token sums its + slots in slot order, so the result does not depend on where the sort placed + each row, on which other tokens share the batch, or on the call. That is what + separates it from the FC2 GEMM's fused finalize, whose reduction order varies + between launches. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode, 1 request without MTP; rows from the sort's block kernel", + spec=_ckpt(1), + ), + Cell(why="1 request at MTP verify width 2", spec=_ckpt(2)), + Cell( + why="1 request at MTP verify width 4; last T the sort places in token order", + spec=_ckpt(4), + ), + Cell(why="decode, 8 requests; rows placed by the sort's atomics", spec=_ckpt(8)), + Cell(why="decode at the 32-request ceiling", spec=_ckpt(32)), + Cell(why="32 requests at MTP verify width 3", spec=_ckpt(96)), + Cell(why="32 requests at MTP verify width 4; largest captured T", spec=_ckpt(128)), + Cell(why="8192 rows, max_num_tokens prefill", spec=_ckpt(8192)), + Cell( + why="8192 rows all on 10 experts; a token's rows in 10 runs of 8192", + spec=_ckpt(8192, dist="hot"), + ), + Cell( + why="510 experts with one row each; most permuted rows are padding", + spec=_ckpt(51, dist="round_robin"), + ), + Cell( + why="FC2 rows spanning 2^-30..2^30; terms 2^60 apart in one sum", + spec=_ckpt(8192, values="wide"), + ), + Cell( + why="tile 256 sort layout, decode, 1 request", + spec=_ckpt(1, tile_tokens_dim=256), + ), + Cell( + why="tile 256 sort layout, 8192-row prefill", + spec=_ckpt(8192, tile_tokens_dim=256), + ), + Cell( + why="slots marked -1 (another rank's expert) skipped; a row of all -1 is 0", + spec=_ckpt(300, non_local_from=256), + ), + ) + + note = """ + Metadata consumed: none. Stateless; allocates nothing, writes `output` in place + on the current stream with one kernel, `moeUnpermuteKernel<__nv_bfloat16, + float, 256>` (`cpp/tensorrt_llm/kernels/moe/cuteDsl/moeUtils.cu`, registered + in `cpp/tensorrt_llm/thop/moe/cuteDsl/cuteDslMoeUtilsOp.cpp`), one thread block + per token. No host synchronization: it captures into a CUDA graph, and replay + on new rows, scales and ids (moe_sort captured with it) is bitwise equal to + eager and to `reference`. + + Call sites: the engine's CuteDSL MoE calls it when the FC2 finalize is not + fused (`run_moe_nvfp4_impl`, `_torch/moe/fused_moe/fused_moe_cute_dsl.py`), and + this target's `_moe_block` calls it in the prefill and the decode forward with + the same four tensors: the FC2 output, the forward's dense `[num_tokens, + hidden]` routed-output buffer, the sort's third output and the router's fp32 + weights. Prefill runs eagerly at up to max_num_tokens (8192) rows; decode runs + captured at requests x MTP verify width, up to 128. + + What it reads: only the `T * top_k` rows `expanded_idx_to_permuted_idx` names, + and the scales. The padding rows of each tile and the tiles past the sort's + `num_non_exiting_tiles` are never read (every cell fills every other row of + `permuted_input` with NaN, and under a graph those rows hold the previous + replay's rows), so the FC2 output's unwritten regions need no memset before + this op. What it writes: every element of `output`, whatever it held, and none + of its inputs. + + Determinism (bitwise): repeated calls; a second moe_sort of the same ids, which + places an expert's rows in another order; a random permutation of all permuted + rows; a prefix of a larger batch, with or without a fresh sort of the prefix + alone; a shuffle of the tokens. The FC2 GEMM's fused finalize + (`cute_dsl_nvfp4_grouped_gemm_finalize_inplace_blackwell`) differs from its + own first call in about half the output elements at decode, while the plain FC2 + GEMM followed by this op does not, and is bitwise equal to the slot-order fp32 + sum of the real FC2 rows. + + Accuracy against the exact sum: the fp32 sum's error is bounded by its terms' + magnitudes, so an element whose terms nearly cancel can miss the float64 sum + rounded once to bf16 by many ulps of the element itself (827 of 20971520 + elements at T = 8192, standard-normal rows and softmax weights, up to 137.5). + + Silent, refused by `is_valid`: a non-contiguous `permuted_input`, `output`, + `expanded_idx_to_permuted_idx` or `topk_scales` (all four are read or written + from `data_ptr()` as dense row-major); an `output` whose dtype is not + `permuted_input`'s (an fp32 output receives the bf16 bit patterns in the first + half of its bytes); 2^31 or more tokens (nothing is written and nothing + raises). + + Data hazards `is_valid` cannot see (values): a permuted index at or past + `permuted_input.size(0)` is read without a bounds check; moe_sort never emits + one. + + Loud: int64 indices; hidden size not a multiple of 8; fp32 rows, or fp16 scales + with bf16 rows; fewer than `T * top_k` permuted rows; a scales or index shape + that disagrees with `output`. + + Accepted, not certified: T = 0 returns and writes nothing; bf16 scales; fp16 + rows; a CPU `output`. + + Certified only at the checkpoint's values: top_k 10, hidden 2560, bf16 rows and + output, fp32 scales, the sort's layouts at tile 128 and 256, every slot local + except in the -1 cell. + + Hard limits (sm_103) + tokens one thread block per token, the count narrowed to int32: + boundary 2^31, `is_valid` refuses it + output elements T * H, offset in int64: exact past 2^31 and 2^32 + permuted rows int32 index, offset in int64; no index passes through fp32 + """ + + def __call__( + self, + permuted_input: torch.Tensor, + output: torch.Tensor, + expanded_idx_to_permuted_idx: torch.Tensor, + topk_scales: torch.Tensor, + ) -> torch.Tensor: + torch.ops.trtllm.moe_unpermute_inplace( + permuted_input=permuted_input, + output=output, + expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx, + topk_scales=topk_scales, + ) + return output + + def reference( + self, + permuted_input: torch.Tensor, + output: torch.Tensor, + expanded_idx_to_permuted_idx: torch.Tensor, + topk_scales: torch.Tensor, + ) -> torch.Tensor: + """The slot-order fp32 sum, one rounding per slot, rounded once at the end. + + Each product of a row and its weight is exact in float64; `_fma_f32` adds + it to the fp32 running sum with one rounding, as a fused multiply-add + does. A slot with a negative index leaves the sum untouched. `output` is + ignored: the result has `permuted_input`'s dtype, which is what the op + writes. + """ + idx = expanded_idx_to_permuted_idx.long() + acc = torch.zeros( + idx.shape[0], permuted_input.shape[1], dtype=torch.float32, device=idx.device + ) + for k in range(idx.shape[1]): + slot = idx[:, k] + rows = permuted_input.index_select(0, slot.clamp_min(0)).double() + prod = rows * topk_scales[:, k : k + 1].double() + acc = torch.where((slot >= 0).unsqueeze(1), _fma_f32(acc, prod), acc) + return acc.to(permuted_input.dtype) + + def is_valid( + self, + permuted_input: torch.Tensor, + output: torch.Tensor, + expanded_idx_to_permuted_idx: torch.Tensor, + topk_scales: torch.Tensor, + ) -> None: + # One thread block per token, the count narrowed to int32: from 2^31 + # tokens nothing is written and no error surfaces. + assert output.size(0) < 2**31, ( + f"output has {output.size(0)} >= 2^31 rows; moe_unpermute_inplace launches one " + "block per token from an int32 count and silently writes nothing" + ) + # The kernel reads and writes all four through data_ptr() as dense + # row-major buffers, and writes `output` with permuted_input's dtype. + for name, t in ( + ("permuted_input", permuted_input), + ("output", output), + ("expanded_idx_to_permuted_idx", expanded_idx_to_permuted_idx), + ("topk_scales", topk_scales), + ): + assert t.is_contiguous(), ( + f"{name} must be contiguous; moe_unpermute_inplace addresses it as a dense " + "row-major buffer and silently combines or writes the wrong elements" + ) + assert output.dtype == permuted_input.dtype, ( + f"output dtype {output.dtype} must be permuted_input's ({permuted_input.dtype}); " + "the op writes permuted_input's element type into it without checking" + ) + + +moe_unpermute_inplace = MoeUnpermuteInplace() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/layer_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/layer_norm.py index 886a2e830c40..84c4deb978a8 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/layer_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/layer_norm.py @@ -53,6 +53,12 @@ class LayerNorm(OpWrapper): ) for rows in (1, 7, 256, 1024, 4096) ) + + ( + Cell( + why="vision norm at the tower call ceiling, 32768 patch rows x 1152 bf16", + spec=dict(shape=(32768, WIDTH), dtype=torch.bfloat16), + ), + ) + tuple( Cell( why=f"{str(dtype)[6:]} activations, {rows} rows", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ngram_token_history.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ngram_token_history.py index 9ad4c4d8230a..50c9d72d4924 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ngram_token_history.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ngram_token_history.py @@ -25,9 +25,10 @@ CKPT_NGRAM_SIZE = 3 CKPT_EOS_TOKEN_ID = 248044 CKPT_VOCAB_SIZE = 248320 -#: Largest `ple_ngram` pool the cache manager was seen to allocate for this -#: checkpoint (10, 35, 366, 368, 396, 399, 432, 461 over twelve runtime-value -#: dumps, batch never above 32). The target hands the whole pool over unsliced. +#: A `ple_ngram` slot count of a full-size engine run. The op addresses a slot by its index +#: alone, so a pool of another size (the pool shares the recurrent pool's slot space, and a +#: bf16 recurrent pool holds more slots in the same memory) takes the same path. The target +#: hands the whole pool over unsliced. ENGINE_SLOTS = 461 @@ -161,6 +162,19 @@ class NgramTokenHistory(OpWrapper): qsl_from="seq_lens", **_CALL_SITE, ), + *( + _cell( + f"MTP verify p={rows}: {decodes} decode{'s' if decodes > 1 else ''} x " + f"{rows} rows, qsl rebuilt", + (rows,) * decodes, + _scattered(decodes), + (True,) * decodes, + ENGINE_SLOTS, + qsl_from="seq_lens", + **_CALL_SITE, + ) + for rows, decodes in ((5, 8), (6, 32), (7, 8), (8, 1), (8, 32)) + ), _cell( "MTP verify mixed: 37-token context + 8 x 3 rows", (37,) + (3,) * 8, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_triton_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_triton_update.py index bbeb9ca56599..a97239b6b4b2 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_triton_update.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_triton_update.py @@ -24,9 +24,10 @@ #: This checkpoint's GDN `[q|k|v]` width; the conv is width 4 over a 3-column pool. CKPT_DIM = 10240 -#: Largest conv-pool slot count the cache manager was seen to allocate for this -#: checkpoint (10 .. 461 over twelve runtime-value dumps). The verify path hands -#: the whole pool over unsliced. +#: A conv-pool slot count of a full-size engine run. The kernel addresses a slot by its +#: index alone, so a pool of another size (a bf16 recurrent pool holds more slots in the +#: same memory, and the conv pool shares its slot space) takes the same path. The verify +#: path hands the whole pool over unsliced. ENGINE_SLOTS = 461 #: The engine's `max_num_sequences`: the snapshot buffer's row count. ENGINE_SNAPSHOT_ROWS = 32 @@ -72,8 +73,8 @@ class CausalConv1dTritonUpdate(OpWrapper): # 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. + # snapshots (rows of the snapshot buffer, or 0 for none), snapshot_steps + # (the buffer's step depth, default seqlen), accepted, launch_dependent_kernels. CELLS: tuple[Cell, ...] = ( tuple( Cell( @@ -87,7 +88,33 @@ class CausalConv1dTritonUpdate(OpWrapper): snapshots=ENGINE_SNAPSHOT_ROWS, ), ) - for batch, seqlen in ((1, 2), (8, 3), (32, 2), (17, 4)) + for batch, seqlen in ( + (1, 2), + (8, 3), + (32, 2), + (17, 4), + # max_draft_len 4..7: seqlen 5..8, up to 256 rows. + (32, 5), + (8, 6), + (17, 7), + (1, 8), + (17, 8), + (32, 8), + ) + ) + + ( + Cell( + why="seqlen 5 into an 8-step snapshot buffer: steps 5..7 untouched", + spec=dict( + dim=CKPT_DIM, + batch=8, + seqlen=5, + pool_slots=ENGINE_SLOTS, + slots=_pinned(8), + snapshots=ENGINE_SNAPSHOT_ROWS, + snapshot_steps=8, + ), + ), ) + tuple( Cell( @@ -215,14 +242,15 @@ class CausalConv1dTritonUpdate(OpWrapper): Metadata consumed: none. Call site: the GDN speculative verify step -- `x` the `[n, 10240, p]` - transposed view of `[n, p, 10240]` (p = runtime_draft_len + 1), the layer's + transposed view of `[n, p, 10240]` (p = runtime_draft_len + 1, 2..8 here), the layer's whole conv pool `[slots, 10240, 3]`, W 4, no bias, `"silu"`, the verify rows' slots, `pad_slot_id` -1, snapshots into `[max_num_sequences, p, 10240, 3]` with `intermediate_state_indices = arange(n)`, `num_accepted_tokens` None. The op advances the pool by all `p` tokens whatever is accepted later; the snapshots are what a rollback has to read. Certified at 128 and 461 pool slots (bitwise identical for the same slot content), snapshot buffers of - `batch` and of 32 rows. + `batch` and of 32 rows, seqlen 2..8 (a buffer deeper than seqlen keeps its + later steps untouched). Raise loudly, so not guarded: width 1 (triton CompilationError); a missing `intermediate_state_indices` with snapshots and slot indices given. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/chunk_gated_delta_rule.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/chunk_gated_delta_rule.py index 703afea315b0..5f87469673ea 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/chunk_gated_delta_rule.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/chunk_gated_delta_rule.py @@ -27,10 +27,9 @@ #: This checkpoint's Gated DeltaNet geometry at tp1. NUM_K_HEADS, NUM_V_HEADS, HEAD_K_DIM, HEAD_V_DIM = 16, 48, 128, 128 -#: Largest slot count the cache manager was seen to allocate the recurrent pool -#: at for this checkpoint (10, 35, 366, 368, 396, 399, 432, 461 over twelve -#: runtime-value dumps; the batch never above 32). The target passes the -#: layer's whole pool unsliced. +#: A recurrent-pool slot count of a full-size engine run. The kernels address a slot by +#: its index alone, so a pool of another size (a bf16 pool holds more slots in the same +#: memory) takes the same path. The target passes the layer's whole pool unsliced. ENGINE_SLOTS = 461 @@ -76,7 +75,8 @@ class ChunkGatedDeltaRule(OpWrapper): # 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] + # | "positional" | None; pool_dtype (fp32 | bf16, the dtype of the indexed or + # positional pool); 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, ...] = ( @@ -101,6 +101,32 @@ class ChunkGatedDeltaRule(OpWrapper): ), ), ) + + tuple( + Cell( + why=f"bf16 pool, 461-slot pool, seqlens {list(s)}", + spec=dict( + seqlens=s, + initial="indexed", + slots=ENGINE_SLOTS, + indices=_engine_slots(len(s)), + pool_dtype=torch.bfloat16, + ), + ) + for s in ((1,), (64,), (200, 37, 91), (1024,), (200, 1, 1, 1, 1)) + ) + + ( + Cell( + why="bf16 pool, 461-slot pool, output= buffer", + spec=dict( + seqlens=(200, 37, 91), + initial="indexed", + slots=ENGINE_SLOTS, + indices=_engine_slots(3), + pool_dtype=torch.bfloat16, + output="buffer", + ), + ), + ) + tuple( Cell( why=f"indexed, seqlens {list(s)}, {len(s) + 3}-slot pool, output= buffer", @@ -150,17 +176,22 @@ class ChunkGatedDeltaRule(OpWrapper): Call site (`_gdn_block`): `(q, k, v, g, beta, None, ssm_states, si_all, True, False, qsl_long, False, None)` after `fused_gdn_post_conv` (q/k - already normalized) -- the layer's whole fp32 pool `[slots, 48, 128, 128]`, - int32 slots, int64 `cu_seqlens`; the MTP-verify context half passes the - `si_p` / `qsl_long_p` prefixes. Certified at small pools and at 461 slots; - at 461 output and state are bitwise identical to a 6-slot pool with the - same slot content. + already normalized) -- the layer's whole pool `[slots, 48, 128, 128]`, fp32 + or bf16 as `mamba_ssm_cache_dtype` allocated it, int32 slots, int64 + `cu_seqlens`; the MTP-verify context half passes the `si_p` / `qsl_long_p` + prefixes, and a mixed batch's one-token decode rows ride in the same call. + Certified at small pools and at 461 slots; at 461 output and state are + bitwise identical to a 6-slot pool with the same slot content. + + Over a bf16 pool the initial state is read as fp32 and the final state is stored + rounded once to bf16, once per call and sequence; the accumulation is unchanged. The first indexed in-place call for an autotune key in a process is as correct as a warm one: the state kernel's autotuner snapshots and restores the indexed state slots - its benchmark launches write (`chunk_delta_h.py`, `_save_autotune_state`). Certified - in a fresh interpreter at this geometry and at 5 / 20 heads. The positional path - never aliased its state. + its benchmark launches write (`chunk_delta_h.py`, `_save_autotune_state`). The key + includes the pool dtype, so a bf16 pool has its own. Certified in a fresh interpreter + at this geometry over an fp32 and over a bf16 pool, 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 @@ -174,8 +205,8 @@ class ChunkGatedDeltaRule(OpWrapper): 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. + Not certified: the FlashInfer sibling, fp16 (accepted), bf16 `g`/`beta`, an + fp16 pool, `K = 256` (accepted, unmeasured), empty `T` or zero-length spans. Hard limits K state kernel asserts K <= 256 (read, not probed) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/cu_seqlens_from_seq_lens.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/cu_seqlens_from_seq_lens.py index 6a10c7d409cd..843e1e9c78a2 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/cu_seqlens_from_seq_lens.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/cu_seqlens_from_seq_lens.py @@ -22,7 +22,7 @@ #: 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). +#: under MTP (2..8: `max_draft_len` 1..7). BATCHES = { "single_prefill": (4096,), "decode_p1": (1,) * 32, @@ -75,6 +75,18 @@ class CuSeqlensFromSeqLens(OpWrapper): layout="prefix", ), ), + *( + Cell( + why=f"MTP verify p={p}, {requests} requests: prefix int32->int64", + spec=dict( + lengths=(p,) * requests, + in_dtype=torch.int32, + out_dtype=torch.int64, + layout="prefix", + ), + ) + for p, requests in ((5, 8), (6, 32), (7, 8), (8, 32)) + ), Cell( why="non-contiguous [::2] seq_lens view", spec=dict( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/dilated_short_conv_paged.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/dilated_short_conv_paged.py index 5bf4c9257981..6ad630e90f1f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/dilated_short_conv_paged.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/dilated_short_conv_paged.py @@ -38,9 +38,10 @@ #: 3, hence the 9-column state the mamba cache manager allocates per slot. CKPT_CHANNELS = 10240 CKPT_DILATION = 3 -#: Largest PLE-conv-pool slot count the cache manager was seen to allocate for -#: this checkpoint (10 .. 461 over twelve runtime-value dumps, batch never -#: above 32). The target hands the whole pool over unsliced. +#: A PLE-conv-pool slot count of a full-size engine run. The kernel addresses a slot by its +#: index alone, so a pool of another size (the pool shares the recurrent pool's slot space, +#: and a bf16 recurrent pool holds more slots in the same memory) takes the same path. The +#: target hands the whole pool over unsliced. ENGINE_SLOTS = 461 @@ -120,6 +121,17 @@ class DilatedShortConvPaged(OpWrapper): x_form="token_major", qsl_dtype=torch.int64, ), + *( + _ckpt( + f"MTP verify p={p} at 461 slots: {n} rows of {p} tokens", + (p,) * n, + pool_slots=ENGINE_SLOTS, + slots=_scattered(n), + x_form="token_major", + qsl_dtype=torch.int64, + ) + for p, n in ((5, 8), (6, 32), (7, 8), (8, 1), (8, 32)) + ), ) + tuple( Cell( @@ -212,7 +224,9 @@ class DilatedShortConvPaged(OpWrapper): 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. + and 461 slots; at 461 bitwise identical to 8 for the same slot content. MTP + verify batches are certified at 3 and at 5..8 tokens per request, up to 32 + requests. 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 @@ -226,8 +240,10 @@ class DilatedShortConvPaged(OpWrapper): Launch-bound on decode (8 rows: 0.461 ms eager vs 0.029 ms of conv; 0.094 ms replayed from a graph). - Not measured: pools above 461 slots, token counts above 8192, dilation - above 3, kernel above 5, `torch.compile`, a pool dtype unlike the model's. + Not measured: pools above 461 slots (a bf16 recurrent state gives the engine + larger ones; a slot is addressed by its index alone), token counts above 8192, + dilation above 3, kernel above 5, `torch.compile`, a pool dtype unlike the + model's. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_cached_replay_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_cached_replay_update.py new file mode 100644 index 000000000000..d131919c8078 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_cached_replay_update.py @@ -0,0 +1,668 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Gated DeltaNet speculative-verify recurrence under the engine's cached-replay +state update: a checkpoint state per slot plus a compact history of the tokens +accepted since it, instead of a per-step state buffer. + +Wraps the in-tree module callable +``tensorrt_llm._torch.modules.fla.cached_replay.fused_recurrent_gated_delta_rule_cached_replay_update`` +(tier-2 sourcing: the recurrence has no op boundary). One launch of the vendored +Triton kernel ``_cached_replay_kernel`` -- plain ``triton.jit``, no autotuner; +the Python body picks the launch configuration from the shapes. +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.modules.fla.cached_replay import ( + fused_recurrent_gated_delta_rule_cached_replay_update as _cached_replay_update, +) + +#: This checkpoint's GDN geometry at tp1. +NUM_K_HEADS, NUM_V_HEADS, HEAD_K_DIM, HEAD_V_DIM = 16, 48, 128, 128 +#: `[q | k | v]` row width of the post-conv activations a verify step packs. +PACKED_QKV_WIDTH = 2 * NUM_K_HEADS * HEAD_K_DIM + NUM_V_HEADS * HEAD_V_DIM # 10240 +#: The cache manager's `replay_history_size`: max(MIN_REPLAY_HISTORY_SIZE=16, T). +HISTORY_SIZE = 16 +#: From this batch size on the engine runs the partitioned replay (work items, no +#: in-kernel commit, one all-layer commit after every layer has run): +#: `CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE` in `modules/fla/cached_replay.py`. +ALL_LAYER_COMMIT_MIN_BATCH = 16 +#: This checkpoint's GDN layer count at tp1 (36 of its 48 layers are linear attention). +NUM_GDN_LAYERS = 36 +#: The recurrent pool's slot count of a full-size engine run. +ENGINE_SLOTS = 461 +#: Layers per slot at which the last of `ENGINE_SLOTS` slots starts past 2^31 elements of +#: the pool (`460 * L * 48 * 128 * 128 > 2**31` from L = 6). +PAST_INT32_LAYERS = 8 + + +def _cell_spec(path, n, t, **spec) -> dict: + return dict(path=path, n=n, t=t, **spec) + + +def _bf16_cell(why, path, n, t, **spec) -> Cell: + return Cell( + why=f"bf16 state {why}", spec=_cell_spec(path, n, t, state_dtype=torch.bfloat16, **spec) + ) + + +class FusedRecurrentGatedDeltaRuleCachedReplayUpdate(OpWrapper): + """One verify step of `N` sequences x `T` tokens through one GDN layer, for + the cached-replay state layout. Per live row (`slot = state_indices[n] >= 0`), + with `p = prev_num_accepted_tokens[slot]`, `b = cache_buf_idx[slot]`: + + 1. the state the step starts from is rebuilt from the checkpoint + `h0 = ssm_states[slot]` (`[HV, V, K]`, K innermost) and the `p` pending + history entries in half `b` -- update vectors `u_j = old_u[slot, b, j]`, + normalized keys `k_j = old_k[slot, b, j]`, cumulative log-decays + `G_j = old_G[slot, b, :, j]` measured from the checkpoint: + `h = exp(G_{p-1}) h0 + sum_j exp(G_{p-1} - G_j) u_j k_j^T`; + 2. the `T` new tokens run the gated delta rule in order (q/k head + `hv // (HV // H)`, so a v-head group shares one q/k head): + `h *= exp(g_t); u_t = beta_t (v_t - h k_t); h += u_t k_t^T; o_t = h (scale q_t)`, + q and k L2-normalized first (`x / sqrt(sum(x^2) + 1e-6)`) when + `use_qk_l2norm_in_kernel`, `scale=None` meaning `K ** -0.5`; with + `A_log` given, `g = -exp(A_log) * softplus(g + dt_bias)` and + `beta = sigmoid(beta)` are applied to the raw projections inside; + 3. the new tokens' `(u_t, k_t, G_t)` are appended to the history. When + `p + T <= history_size` they go to half `b` at `p .. p+T-1` with `G` + still cumulative from the checkpoint. Otherwise (a *wrap*) they go to + half `1 - b` at `0 .. T-1` with `G` restarted from zero, and the + checkpoint must advance to the state after the `p` pending tokens: + without work items the kernel writes that state into `ssm_states[slot]` + itself; with `replay_work_items` (`use_all_layer_commit=True`) it does + not, and the caller owes one `commit_gdn_cached_replay_history_layers` + over every layer before the next step. + + Nothing about acceptance happens here: the caller advances `p`/`b` after + sampling (`p <- accepted` and `b <- 1 - b` on a wrap, else `p += accepted`, + `accepted` counting the bonus token), so rejected tokens' history entries + are simply never read. + + Fusion boundary: L2 norms, scale, the gating of the raw `a`/`b` + projections, the history rebuild, the recurrence, the history append and + (small batches) the checkpoint commit are inside. The conv, the projections, + the all-layer commit and the `p`/`b` bookkeeping are not. + + Returns `[N, T, HV, V]` in `output`'s dtype, else `q`'s; it *is* `output` + when one is given. `ssm_states`, `old_u`, `old_k`, `old_G` are written in + place and not returned. Rows with a negative slot are skipped entirely: + nothing is read or written for them, their output rows included. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: path ("small": no work items, the kernel commits a wrap itself | + # "all_layer": work items + use_all_layer_commit, the test then runs the + # manager's all-layer commit), n sequences x t tokens per verify step, + # steps (consecutive verify steps, each followed by an acceptance per row), + # accept ("random": 1..t per row, the default | "full": t every step, so from + # a fresh slot at t=8 the second step appends at `pnat + t == 16` without + # wrapping and every second step after wraps | "one": 1 every step), + # state_dtype (the checkpoint pool's dtype: fp32 | bf16, as + # `mamba_ssm_cache_dtype` allocates it; fp32 by default), + # layers (each with its own activations, gates + # and history; one V2-style pool interleaving the layers per slot), slots + # (the pool's slot count), start ("warm": random pending histories in both + # halves, a wrap and an append forced at step 0 | "fresh": pnat 0 and + # zeroed history, as the manager leaves a slot after prefill), pad (rows + # given slot -1), edge_slots (rows 0 and n-1 on the pool's last and first + # slot), output (None | "prefix" of a longer buffer | "suffix": its rows + # after 5 context tokens), packed (False: the q/k/v views over the packed + # rows without packed_qkv, as a mixed batch passes them), pdl + # (launch_with_pdl, default False, the form the target calls: PDL is a constexpr of the + # kernel, so True is a separate compile), ab ("split": separate [n, t, 48] + # projections | "strided": the column halves of one [n*t, 96] projection), + # dup (rows given the SAME slot, as CUDA-graph padding repeats one dummy slot; + # only the other rows and slots are held to the reference). + CELLS: tuple[Cell, ...] = ( + Cell(why="small n=1 T=2: one sequence, MTP-1", spec=_cell_spec("small", 1, 2, steps=12)), + Cell(why="small n=4 T=3", spec=_cell_spec("small", 4, 3)), + Cell( + why="small n=8 T=4 L=2 S=461: slots 0 and 460 of a 461-slot axis", + spec=_cell_spec("small", 8, 4, slots=ENGINE_SLOTS, edge_slots=True), + ), + Cell( + why="small n=15 T=2: largest batch before the split", + spec=_cell_spec("small", 15, 2, steps=12), + ), + Cell( + why="small n=8 T=4 fresh: first verify after prefill", + spec=_cell_spec("small", 8, 4, start="fresh", steps=10), + ), + Cell( + why="small n=6 T=3 pad: rows 1, 4 at slot -1, prefix output", + spec=_cell_spec("small", 6, 3, pad=(1, 4), output="prefix"), + ), + Cell( + why="small n=8 T=3 dup: rows 5..7 share one padding slot", + spec=_cell_spec("small", 8, 3, dup=(5, 6, 7), steps=6), + ), + Cell( + why="small n=4 T=2 unpacked: q/k/v views alone", + spec=_cell_spec("small", 4, 2, packed=False, steps=12), + ), + Cell( + why="small n=4 T=4 PDL on: launch_with_pdl=True, the kernel's other compile", + spec=_cell_spec("small", 4, 4, pdl=True), + ), + Cell( + why="small n=5 T=3 strided a/b: halves of [n*t, 96]", + spec=_cell_spec("small", 5, 3, ab="strided"), + ), + Cell( + why="all_layer n=16 T=4: the partition threshold", spec=_cell_spec("all_layer", 16, 4) + ), + Cell( + why="all_layer n=24 T=3 L=2 S=461: slots 0 and 460 of a 461-slot axis", + spec=_cell_spec("all_layer", 24, 3, slots=ENGINE_SLOTS, edge_slots=True), + ), + Cell( + why="all_layer n=32 T=4: the envelope's 32 requests", + spec=_cell_spec("all_layer", 32, 4), + ), + Cell( + why="all_layer n=32 T=3 L=36 S=40: every GDN layer, one commit", + spec=_cell_spec("all_layer", 32, 3, layers=NUM_GDN_LAYERS, slots=40, steps=6), + ), + Cell( + why="all_layer n=16 T=3 L=8 S=461: slot 460's offset is past 2^31 elements", + spec=_cell_spec( + "all_layer", + 16, + 3, + layers=PAST_INT32_LAYERS, + slots=ENGINE_SLOTS, + edge_slots=True, + steps=3, + ), + ), + Cell( + why="all_layer n=20 T=4 fresh: after prefill", + spec=_cell_spec("all_layer", 20, 4, start="fresh", steps=10), + ), + Cell( + why="all_layer n=18 T=3 pad: rows 0, 17 at slot -1", + spec=_cell_spec("all_layer", 18, 3, pad=(0, 17), output="prefix"), + ), + Cell( + why="all_layer n=16 T=3 mixed-batch form: views, offset out", + spec=_cell_spec("all_layer", 16, 3, packed=False, output="suffix"), + ), + # ---- MTP draft lengths 4..7: T = 5..8, the kernel's ceiling ---- + Cell(why="small n=1 T=5: one sequence, MTP-4", spec=_cell_spec("small", 1, 5, steps=10)), + Cell(why="small n=8 T=6: MTP-5", spec=_cell_spec("small", 8, 6)), + Cell( + why="small n=15 T=7: MTP-6, largest batch before the split", + spec=_cell_spec("small", 15, 7), + ), + Cell(why="small n=1 T=8: one sequence, MTP-7", spec=_cell_spec("small", 1, 8, steps=10)), + Cell(why="small n=8 T=8: MTP-7", spec=_cell_spec("small", 8, 8)), + Cell( + why="small n=15 T=8: MTP-7, largest batch before the split", + spec=_cell_spec("small", 15, 8), + ), + Cell( + why="small n=8 T=8 S=461: slots 0 and 460 of a 461-slot axis", + spec=_cell_spec("small", 8, 8, slots=ENGINE_SLOTS, edge_slots=True), + ), + Cell( + why="small n=4 T=8 full accept: pnat+T==16 appends, then wraps", + spec=_cell_spec("small", 4, 8, start="fresh", accept="full", steps=8), + ), + Cell( + why="small n=8 T=8 one accepted: a=1 every step", + spec=_cell_spec("small", 8, 8, accept="one", steps=12), + ), + Cell( + why="small n=6 T=8 mixed-batch form: views, offset out, pad", + spec=_cell_spec("small", 6, 8, packed=False, output="suffix", pad=(2,)), + ), + Cell(why="all_layer n=16 T=5: MTP-4 at the threshold", spec=_cell_spec("all_layer", 16, 5)), + Cell( + why="all_layer n=24 T=6 S=461: MTP-5, slots 0 and 460", + spec=_cell_spec("all_layer", 24, 6, slots=ENGINE_SLOTS, edge_slots=True), + ), + Cell(why="all_layer n=32 T=7: MTP-6 at 32 requests", spec=_cell_spec("all_layer", 32, 7)), + Cell(why="all_layer n=16 T=8: MTP-7 at the threshold", spec=_cell_spec("all_layer", 16, 8)), + Cell(why="all_layer n=32 T=8: MTP-7, 256 rows", spec=_cell_spec("all_layer", 32, 8)), + Cell( + why="all_layer n=32 T=8 L=36 S=40: every GDN layer, one commit", + spec=_cell_spec("all_layer", 32, 8, layers=NUM_GDN_LAYERS, slots=40, steps=6), + ), + Cell( + why="all_layer n=32 T=8 full accept: commit every second step", + spec=_cell_spec("all_layer", 32, 8, start="fresh", accept="full", steps=8), + ), + Cell( + why="all_layer n=16 T=8 mixed-batch form: views, offset out", + spec=_cell_spec("all_layer", 16, 8, packed=False, output="suffix"), + ), + Cell( + why="all_layer n=20 T=8 fresh pad: rows 0, 19 at slot -1", + spec=_cell_spec("all_layer", 20, 8, start="fresh", pad=(0, 19), output="prefix"), + ), + # ---- a bf16 checkpoint pool (`mamba_ssm_cache_dtype=bfloat16`) ---- + _bf16_cell("small n=1 T=2: one sequence, MTP-1", "small", 1, 2, steps=12), + _bf16_cell("small n=8 T=4", "small", 8, 4), + _bf16_cell("small n=15 T=3: largest batch before the split", "small", 15, 3), + _bf16_cell("small n=8 T=4 fresh: first verify after prefill", "small", 8, 4, start="fresh"), + _bf16_cell( + "small n=6 T=3 pad: rows 1, 4 at slot -1, prefix output", + "small", + 6, + 3, + pad=(1, 4), + output="prefix", + ), + _bf16_cell( + "small n=8 T=4 S=461: slots 0 and 460 of a 461-slot axis", + "small", + 8, + 4, + slots=ENGINE_SLOTS, + edge_slots=True, + ), + _bf16_cell("small n=1 T=5: MTP-4", "small", 1, 5, steps=10), + _bf16_cell("small n=1 T=8: MTP-7, one sequence", "small", 1, 8, steps=10), + _bf16_cell("small n=8 T=8: MTP-7", "small", 8, 8), + _bf16_cell( + "small n=4 T=8 full accept: pnat+T==16 appends, then wraps", + "small", + 4, + 8, + start="fresh", + accept="full", + steps=8, + ), + _bf16_cell("all_layer n=16 T=4: the partition threshold", "all_layer", 16, 4), + _bf16_cell("all_layer n=24 T=6: MTP-5", "all_layer", 24, 6), + _bf16_cell("all_layer n=32 T=8: MTP-7, 256 rows", "all_layer", 32, 8), + _bf16_cell( + "all_layer n=32 T=3 L=36 S=40: every GDN layer, one commit", + "all_layer", + 32, + 3, + layers=NUM_GDN_LAYERS, + slots=40, + steps=6, + ), + _bf16_cell( + "all_layer n=32 T=8 full accept: commit every second step", + "all_layer", + 32, + 8, + start="fresh", + accept="full", + steps=8, + ), + _bf16_cell( + "all_layer n=16 T=8 mixed-batch form: views, offset out", + "all_layer", + 16, + 8, + packed=False, + output="suffix", + ), + ) + + note = """ + Call sites. The engine's GDN mixer (`modules/mamba/gdn_mixer.py`) is the + reference caller: the target-verify step of a generation-only batch, and the + decode rows of a mixed context + verify batch -- `N = num_decodes`, `T = + runtime_draft_len + 1` (2..8 here: max_draft_len 1..7). A plain decode step never reaches it. q/k/v + are the `[N, T, 16|48, 128]` views of the post-conv `[N*T, 10240]` rows, plus + `packed_qkv` = those rows for a generation-only batch (a mixed batch passes the + views alone and the op copies them contiguous); `g=a`, `beta=b` are the raw + bf16 `[N, T, 48]` projections, with the layer's fp32 `[48]` `A_log` and + `dt_bias`; `ssm_states` is the layer's pool view `[S, 48, 128, 128]`, fp32 or bf16 as + `mamba_ssm_cache_dtype` allocated it (V2 interleaves the layers inside each slot, so + the slot stride is a multiple of one state); `state_indices` the decode rows' int32 slots; `old_u, old_k, old_G, + old_beta` the layer's `old_x, old_B, old_dt, old_dA_cumsum` from + `mamba_layer_cache(layer)`; `cache_buf_idx`, `prev_num_accepted_tokens`, + `history_size` from `get_replay_state_update_metadata()` (16 here); + `use_qk_l2norm_in_kernel=True`; and exactly when `N >= 16` and the manager's + `use_gdn_cached_replay_all_layer_commit` holds, `replay_work_items = + mamba_metadata.replay_work_items[:N]`, `n_writes = mamba_metadata.replay_n_writes`, + `use_all_layer_commit=True`. That condition has to be the manager's own: + `update_mamba_states` runs the all-layer commit iff it holds, so work items + passed when it does not leave every wrap uncommitted, and work items withheld + when it does commit twice -- both silent at the op and far off the exact + recurrence one step later. + + Metadata consumed: `state_indices` (one slot per row), and per slot + `prev_num_accepted_tokens` and `cache_buf_idx`. `replay_work_items` / + `n_writes` are only checked (shape, dtype); the kernel never reads them, their + presence switches the in-kernel commit off. `old_beta` is never read or written. + + Preconditions not checkable from metadata: + * rows on distinct slots, except padding: two rows on one slot both append to the + same history positions and the appended u is a per-element mixture of the two + rows, so that slot's state and those rows' outputs are unspecified. The rows of + CUDA-graph padding repeat one shared dummy slot (the manager's `_padding_slot`); + the rows on unique slots are exact whatever the padding does, and the dummy slot + is nobody's state; + * `prev_num_accepted_tokens` in `[0, history_size]`, `cache_buf_idx` in `{0, 1}`, + and the pending entries written by earlier calls of this op. A slot fresh from + prefill starts at 0 pending (the manager zeroes the count, the half and the + history); + * after every step, before the next: with work items, + `commit_gdn_cached_replay_history_layers` over every local GDN layer; then the + advance (`_advance_replay_state` in `pyexecutor/kv_cache/mamba_cache_manager.py`) + with the accepted counts including the bonus token. Advancing as if every + draft were accepted is silent at the op and wrong from the next step on; + * the conv state is not this op's: under replay the conv still writes its + per-step window and the manager promotes it. + + The launch is one V tile, `num_warps=4`, `num_stages=1`, grid `(1, N * 48)`: the + tuned mappings need HV = 2H or 4H and never engage at HV/H = 3. The checkpoint's dtype + is the kernel's `STATE_FP32` specialization and nothing else below N = 128: an fp32 + checkpoint enters the history rebuild's tensor-core dots as a bf16 hi/lo pair, a bf16 + one as stored, and every commit is computed in fp32 and rounded once to the pool's + dtype. From N = 128 a bf16 state alone also turns on the L2-streaming mapping + (`main_num_stages=3`, a 5-stage commit), which no cell drives. `FLA_USE_FAST_OPS=1` at + import switches the kernel's `exp` to the approximate one; the cells run with it unset. + + Bit for bit the same call: the in-kernel commit and the partitioned call plus + the all-layer commit; `packed_qkv` and the separate q/k/v (whose contents are + never read when `packed_qkv` is given); `launch_with_pdl` on and off; a second + run; CUDA-graph replay on new data (the test captures every T = 2..8 at N = 1, 8, 15, 16 + and 32). + + T = 5..8 share one compiled token tile (`BT = next_pow2(T) = 8`). At + T = 8 and history 16 a slot whose every draft is accepted appends rows 8..15 at + `pnat + T == 16` without wrapping and wraps on the next step, so it commits every second + step; the `accept="full"` cells drive that boundary (in-kernel and all-layer commit), the + `accept="one"` cell the slowest fill. + + Loud refusals: work items without `use_all_layer_commit` (ValueError) and the + reverse (ValueError); one of `replay_work_items` / `n_writes` alone + (ValueError); a strided `output`; a strided head axis of `g`; `history_size < + T` (a bare assert); and `history_size <= 8`, which does not compile -- the + history tile is `next_pow2(history_size)` wide and feeds a `tl.dot` whose + reduction must be 16, so only 9..16 work (the engine's is 16). + + `is_valid` guards what is accepted and silently wrong: `state_indices` shorter + than N (the trailing rows read past it and append to a slot nobody named); + history buffers whose shape disagrees with `history_size` or the head geometry + (the kernel addresses them by `history_size`, so 16-deep buffers called at 12 + read other slots' entries); history buffers, `prev_num_accepted_tokens` or + `cache_buf_idx` covering fewer slots than `ssm_states` (a slot of the pool past + their end is read and written out of bounds); a strided `prev_num_accepted_tokens` + or `cache_buf_idx` (read at element `slot`). A strided `state_indices` is + harmless: the entry's input guard makes it contiguous. + + Not certified: pre-gated `g`/`beta` (`A_log=None`); `use_qk_l2norm_in_kernel= + False`; an `ssm_states` other than fp32 or bf16, or a bf16 one at N >= 128; HV/H other + than 3, K or V other than 128; T outside 2..8 (T = 9 raises: the kernel's bare + `assert T <= 8`), at either `ssm_states` dtype; `history_size` other than 16; `old_u` / + `old_k` other than bf16, `old_G` other than fp32; a `scale`; any launch-knob + override (`block_v`, `num_warps`, `use_l2_*`, `main_num_stages`); an output + dtype other than bf16; int64 `state_indices`; N above 32. + """ + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + ssm_states: torch.Tensor, + state_indices: torch.Tensor, + old_u: torch.Tensor, + old_k: torch.Tensor, + old_G: torch.Tensor, + old_beta: torch.Tensor, + cache_buf_idx: torch.Tensor, + prev_num_accepted_tokens: torch.Tensor, + history_size: int, + scale: Optional[float] = None, + use_qk_l2norm_in_kernel: bool = False, + A_log: Optional[torch.Tensor] = None, + dt_bias: Optional[torch.Tensor] = None, + launch_with_pdl: bool = False, + replay_work_items: Optional[torch.Tensor] = None, + n_writes: Optional[torch.Tensor] = None, + block_v: Optional[int] = None, + num_warps: Optional[int] = None, + use_l2_state_cache: Optional[bool] = None, + use_l2_streaming_inputs: Optional[bool] = None, + use_l2_shared_inputs: Optional[bool] = None, + main_num_stages: Optional[int] = None, + packed_qkv: Optional[torch.Tensor] = None, + use_all_layer_commit: bool = False, + output: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return _cached_replay_update( + q=q, + k=k, + v=v, + g=g, + beta=beta, + ssm_states=ssm_states, + state_indices=state_indices, + old_u=old_u, + old_k=old_k, + old_G=old_G, + old_beta=old_beta, + cache_buf_idx=cache_buf_idx, + prev_num_accepted_tokens=prev_num_accepted_tokens, + history_size=history_size, + scale=scale, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + A_log=A_log, + dt_bias=dt_bias, + launch_with_pdl=launch_with_pdl, + replay_work_items=replay_work_items, + n_writes=n_writes, + block_v=block_v, + num_warps=num_warps, + use_l2_state_cache=use_l2_state_cache, + use_l2_streaming_inputs=use_l2_streaming_inputs, + use_l2_shared_inputs=use_l2_shared_inputs, + main_num_stages=main_num_stages, + packed_qkv=packed_qkv, + use_all_layer_commit=use_all_layer_commit, + output=output, + ) + + def reference( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + ssm_states: torch.Tensor, + state_indices: torch.Tensor, + old_u: torch.Tensor, + old_k: torch.Tensor, + old_G: torch.Tensor, + old_beta: torch.Tensor, + cache_buf_idx: torch.Tensor, + prev_num_accepted_tokens: torch.Tensor, + history_size: int, + scale: Optional[float] = None, + use_qk_l2norm_in_kernel: bool = False, + A_log: Optional[torch.Tensor] = None, + dt_bias: Optional[torch.Tensor] = None, + launch_with_pdl: bool = False, + replay_work_items: Optional[torch.Tensor] = None, + n_writes: Optional[torch.Tensor] = None, + block_v: Optional[int] = None, + num_warps: Optional[int] = None, + use_l2_state_cache: Optional[bool] = None, + use_l2_streaming_inputs: Optional[bool] = None, + use_l2_shared_inputs: Optional[bool] = None, + main_num_stages: Optional[int] = None, + packed_qkv: Optional[torch.Tensor] = None, + use_all_layer_commit: bool = False, + output: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """`(out, ssm_states, old_u, old_k, old_G)` as the call leaves them; + mutates nothing. The state is rebuilt from the checkpoint and history + in fp32 and the new tokens are stepped one at a time in fp32 -- no + matmul, so no TF32 -- and every store is rounded once to its buffer's + dtype. The history is read at whatever dtype and depth its buffers + have, so the same function driven on fp32 history deeper than any run + ever fills is the exact token-by-token recurrence. + + Rows with a negative slot keep `output`'s prior content, or are NaN + (unspecified) when no `output` is given. `old_beta`, `launch_with_pdl`, + `n_writes` and the launch-tuning knobs are ignored; `replay_work_items` + only decides whether a wrap commits the checkpoint here.""" + N, T, H, K = k.shape + HV, V = v.shape[2], v.shape[3] + rep = HV // H + dev = v.device + scale = K**-0.5 if scale is None else scale + if packed_qkv is not None: + rows = packed_qkv.view(N, T, -1) + q = rows[..., : H * K].view(N, T, H, K) + k = rows[..., H * K : 2 * H * K].view(N, T, H, K) + v = rows[..., 2 * H * K :].view(N, T, HV, V) + qf, kf, vf = q.float(), k.float(), v.float() + if use_qk_l2norm_in_kernel: + qf = qf / torch.sqrt(qf.square().sum(-1, keepdim=True) + 1e-6) + kf = kf / torch.sqrt(kf.square().sum(-1, keepdim=True) + 1e-6) + qf = (qf * scale).repeat_interleave(rep, dim=2) # [N, T, HV, K] + kv = kf.repeat_interleave(rep, dim=2) # [N, T, HV, K] + if A_log is not None: + gf = -torch.exp(A_log.float()) * F.softplus(g.float() + dt_bias.float()) + bf = torch.sigmoid(beta.float()) + else: + gf, bf = g.float(), beta.float() + + out_dtype = q.dtype if output is None else output.dtype + out = ( + torch.full((N, T, HV, V), float("nan"), device=dev) + if output is None + else output.float().clone() + ) + pool, u_buf, k_buf, G_buf = (x.clone() for x in (ssm_states, old_u, old_k, old_G)) + live = (state_indices[:N] >= 0).nonzero().flatten() + if live.numel() == 0: + return out.to(out_dtype), pool, u_buf, k_buf, G_buf + slot = state_indices[:N][live].long() + p = prev_num_accepted_tokens[slot].long() + b = cache_buf_idx[slot].long() + M, depth = live.numel(), old_u.shape[2] + + # 1. the pending state: checkpoint plus the p history entries of half b + pending = torch.arange(depth, device=dev)[None, :] < p[:, None] # [M, depth] + Gh = torch.where(pending[:, None, :], old_G[slot, b].float(), 0.0) # [M, HV, depth] + Uh = torch.where(pending[:, :, None, None], old_u[slot, b].float(), 0.0) + Kh = torch.where(pending[:, :, None, None], old_k[slot, b].float(), 0.0) + Kh = Kh.repeat_interleave(rep, dim=2) # [M, depth, HV, K] + last = (p - 1).clamp(min=0)[:, None, None].expand(M, HV, 1) + g_start = torch.where(p[:, None] > 0, Gh.gather(2, last).squeeze(2), 0.0) # [M, HV] + weight = torch.where(pending[:, None, :], torch.exp(g_start[:, :, None] - Gh), 0.0) + h = torch.exp(g_start)[:, :, None, None] * ssm_states[slot].float() # [M, HV, V, K] + for j in range(depth): + h = h + weight[:, :, j, None, None] * Uh[:, j, :, :, None] * Kh[:, j, :, None, :] + checkpoint = h + + # 2. the new tokens, in order + gl, bl = gf[live], bf[live] # [M, T, HV] + ql, kl, vl = qf[live], kv[live], vf[live] + G_local = torch.zeros(M, HV, device=dev) + u_new, G_new = [], [] + for t in range(T): + h = h * torch.exp(gl[:, t])[:, :, None, None] + u = bl[:, t, :, None] * (vl[:, t] - (h * kl[:, t, :, None, :]).sum(-1)) # [M, HV, V] + h = h + u[:, :, :, None] * kl[:, t, :, None, :] + out[live, t] = (h * ql[:, t, :, None, :]).sum(-1) + G_local = G_local + gl[:, t] + u_new.append(u) + G_new.append(G_local.clone()) + + # 3. the history append, and the checkpoint commit on a wrap + wrap = p + T > history_size + w_buf = torch.where(wrap, 1 - b, b) + w_off = torch.where(wrap, 0, p) + for t in range(T): + pos = w_off + t + u_buf[slot, w_buf, pos] = u_new[t].to(u_buf.dtype) + k_buf[slot, w_buf, pos] = kf[live, t].to(k_buf.dtype) + G_store = torch.where(wrap[:, None], G_new[t], g_start + G_new[t]) + G_buf[slot, w_buf, :, pos] = G_store.to(G_buf.dtype) + if replay_work_items is None and bool(wrap.any()): + pool[slot[wrap]] = checkpoint[wrap].to(pool.dtype) + return out.to(out_dtype), pool, u_buf, k_buf, G_buf + + def is_valid( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + ssm_states: torch.Tensor, + state_indices: torch.Tensor, + old_u: torch.Tensor, + old_k: torch.Tensor, + old_G: torch.Tensor, + old_beta: torch.Tensor, + cache_buf_idx: torch.Tensor, + prev_num_accepted_tokens: torch.Tensor, + history_size: int, + scale: Optional[float] = None, + use_qk_l2norm_in_kernel: bool = False, + A_log: Optional[torch.Tensor] = None, + dt_bias: Optional[torch.Tensor] = None, + launch_with_pdl: bool = False, + replay_work_items: Optional[torch.Tensor] = None, + n_writes: Optional[torch.Tensor] = None, + block_v: Optional[int] = None, + num_warps: Optional[int] = None, + use_l2_state_cache: Optional[bool] = None, + use_l2_streaming_inputs: Optional[bool] = None, + use_l2_shared_inputs: Optional[bool] = None, + main_num_stages: Optional[int] = None, + packed_qkv: Optional[torch.Tensor] = None, + use_all_layer_commit: bool = False, + output: Optional[torch.Tensor] = None, + ) -> None: + N, _, H, K = k.shape + HV, V = v.shape[2], v.shape[3] + assert state_indices.dim() == 1 and state_indices.shape[0] >= N, ( + f"state_indices has shape {tuple(state_indices.shape)} for {N} sequences -- the " + "kernel reads one slot per sequence with no bound, so a shorter tensor hands the " + "trailing rows whatever follows it in memory, and they write that slot's history" + ) + pool_slots, slots = ssm_states.shape[0], old_u.shape[0] + assert slots >= pool_slots, ( + f"the history buffers have {slots} slots for a {pool_slots}-slot pool -- a slot " + "past the history's end is read and written out of bounds" + ) + for name, tensor in ( + ("prev_num_accepted_tokens", prev_num_accepted_tokens), + ("cache_buf_idx", cache_buf_idx), + ): + assert tensor.dim() == 1 and tensor.is_contiguous() and tensor.shape[0] >= pool_slots, ( + f"{name} has shape {tuple(tensor.shape)} and stride {tensor.stride()}; the kernel " + f"reads one element per slot of the {pool_slots}-slot pool, at the slot's index" + ) + expected = { + "old_u": (old_u, (slots, 2, history_size, HV, V)), + "old_k": (old_k, (slots, 2, history_size, H, K)), + "old_G": (old_G, (slots, 2, HV, history_size)), + } + for name, (tensor, shape) in expected.items(): + assert tuple(tensor.shape) == shape, ( + f"{name} has shape {tuple(tensor.shape)}; history_size={history_size} and this " + f"geometry need {shape} -- the kernel addresses the history by history_size and " + "the head geometry, never by the buffer's own shape" + ) + + +fused_recurrent_gated_delta_rule_cached_replay_update = ( + FusedRecurrentGatedDeltaRuleCachedReplayUpdate() +) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_update.py index a290b44717ab..13209d0dd5b3 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_update.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_update.py @@ -27,8 +27,8 @@ #: 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 +#: A recurrent-pool slot count of a full-size engine run (a bf16 pool holds more slots in +#: the same memory). 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. @@ -47,7 +47,7 @@ class FusedRecurrentGatedDeltaRuleUpdate(OpWrapper): """Advance each sequence's `[K, V]` recurrent state by *all* its tokens, in order. Per token and v-head `hv` (q/k head `hv // (HV // H)`): - q, k = q / (||q|| + 1e-6), k / (||k|| + 1e-6) # if use_qk_l2norm_in_kernel + q, k = q / sqrt(q.q + 1e-6), k / sqrt(k.k + 1e-6) # if use_qk_l2norm_in_kernel h *= exp(g); u = (v - h^T k) * beta; h += k u^T; o = h^T (scale * q) in fp32 throughout. `h` starts at `initial_state_source[idx]` (`[S, HV, V, @@ -93,6 +93,32 @@ class FusedRecurrentGatedDeltaRuleUpdate(OpWrapper): ) for n, p in ((1, 2), (8, 2), (16, 2), (32, 2), (8, 3), (8, 4)) ) + + tuple( + _verify( + n, + p, + f"bf16 state: target verify n={n} p={p}, n-row gather, 32-row buffer", + pool_slots=n, + cache_rows=ENGINE_VERIFY_ROWS, + beta="bf16", + qkv="contiguous", + pool_dtype=torch.bfloat16, + cache_dtype=torch.bfloat16, + ) + for n, p in ((1, 2), (8, 2), (16, 2), (32, 2), (8, 3), (8, 4)) + ) + + ( + _verify( + 4, + 2, + "bf16 state: verify, -1-padded rows naming 31 and 0", + pool_slots=ENGINE_VERIFY_ROWS, + cache_rows=ENGINE_VERIFY_ROWS, + slots=(31, -1, 0, -1), + pool_dtype=torch.bfloat16, + cache_dtype=torch.bfloat16, + ), + ) + ( _verify( 4, @@ -168,16 +194,16 @@ class FusedRecurrentGatedDeltaRuleUpdate(OpWrapper): note = """ Metadata consumed: none; every piece of state is an argument. - Call site (`_gdn_verify`), all positional: `(q, k, v, g, beta, src, - int_idx, None, True, None, True, False, int_ssm, p, None)` -- q/k/v the - contiguous `[n, p, 16|48, 128]` reshapes of `pack_gdn_decode_qkv`'s - outputs (so the in-kernel L2 norm is required), `g` fp32 from - `fused_gdn_gating`, `beta = sigmoid(b)` bf16, `src` an `[n, 48, 128, 128]` - fp32 gather of the live slots, `int_idx = intermediate_state_indices[:n]` - (int32 arange), `int_ssm` the layer's `[32, p, 48, 128, 128]` per-step - buffer. The buffer's rows are the batch rows `update_mamba_states` reads, - not slot ids, and its step axis is `tokens_per_gen_step`: `p` (as - `cache_steps`) must equal it. + Call site (`_gdn_verify`, the legacy verify path of `TRTLLM_USE_GDN_REPLAY=0`, `p` up + to 4), all positional: `(q, k, v, g, beta, src, int_idx, None, True, None, True, False, + int_ssm, p, None)` -- q/k/v the contiguous `[n, p, 16|48, 128]` reshapes of + `pack_gdn_decode_qkv`'s outputs (so the in-kernel L2 norm is required), `g` fp32 from + `fused_gdn_gating`, `beta = sigmoid(b)` bf16, `src` an `[n, 48, 128, 128]` gather of + the live slots and `int_ssm` the layer's `[32, p, 48, 128, 128]` per-step buffer, both + at the pool's dtype (fp32 or bf16, as `mamba_ssm_cache_dtype` allocated it; the buffer's + states are rounded once to it), `int_idx = intermediate_state_indices[:n]` (int32 + arange). The buffer's rows are the batch rows `update_mamba_states` reads, not slot + ids, and its step axis is `tokens_per_gen_step`: `p` (as `cache_steps`) must equal it. Preconditions not checkable from metadata: `HV % H == 0` (q/k heads are repeat_interleaved, not tiled); indices in `[0, S)` and `[0, S')` or @@ -295,8 +321,8 @@ def reference( q_t = q[b, t].float().repeat_interleave(rep, dim=0) # [HV, K] k_t = k[b, t].float().repeat_interleave(rep, dim=0) if use_qk_l2norm_in_kernel: - q_t = q_t / (q_t.square().sum(-1, keepdim=True).sqrt() + 1e-6) - k_t = k_t / (k_t.square().sum(-1, keepdim=True).sqrt() + 1e-6) + q_t = q_t / torch.sqrt(q_t.square().sum(-1, keepdim=True) + 1e-6) + k_t = k_t / torch.sqrt(k_t.square().sum(-1, keepdim=True) + 1e-6) q_t = q_t * scale h = h * torch.exp(g[b, t].float())[:, None, None] u = v[b, t].float() - (h * k_t[:, :, None]).sum(1) # [HV, V] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_sigmoid_gating_delta_rule_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_sigmoid_gating_delta_rule_update.py index 12dd485c4bfe..c822cc138903 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_sigmoid_gating_delta_rule_update.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_sigmoid_gating_delta_rule_update.py @@ -22,9 +22,9 @@ #: The pool width the bulk of the cells run at, from the old test: a width the #: engine never allocates, kept because its cells were certified there. SLOTS = 24 -#: Largest slot count the cache manager was seen to allocate the recurrent pool -#: at for this checkpoint (10 .. 461 over twelve runtime-value dumps, batch never -#: above 32). The target passes the layer's whole pool unsliced. +#: A recurrent-pool slot count of a full-size engine run. The kernels address a slot by +#: its index alone, so a pool of another size (a bf16 pool holds more slots in the same +#: memory) takes the same path. The target passes the layer's whole pool unsliced. ENGINE_SLOTS = 461 @@ -42,7 +42,7 @@ class FusedSigmoidGatingDeltaRuleUpdate(OpWrapper): (stored `[V, K]`, K innermost), per token, all in fp32: g = -exp(A_log) * softplus_b(a + dt_bias) beta = sigmoid(b) - q, k = q / (|q| + 1e-6), k / (|k| + 1e-6) (only with use_qk_l2norm_in_kernel) + q, k = q / sqrt(q.q + 1e-6), k / sqrt(k.k + 1e-6) (only with use_qk_l2norm_in_kernel) h *= exp(g); u = (v - h^T k) * beta; h += k u^T; o = h^T (scale * q) `softplus_b(x) = log(1 + exp(beta_s x)) / beta_s` below `softplus_threshold` @@ -67,7 +67,10 @@ class FusedSigmoidGatingDeltaRuleUpdate(OpWrapper): # 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). + # pool[1] of a [3, slots, ...] cache | "batched": cu_seqlens None, + # contiguous copies | "rows": cu_seqlens None, one token per row, q/k/v + # [N, 1, heads, 128] views of the packed [N, 10240] block whose token + # stride is the row stride, and a/b [N, 1, 48] -- the target's bf16 call). CELLS: tuple[Cell, ...] = ( tuple( Cell( @@ -93,6 +96,43 @@ class FusedSigmoidGatingDeltaRuleUpdate(OpWrapper): ), ), ) + + tuple( + Cell( + why=f"target bf16 decode call, rows form N={n}, 461-slot pool", + spec=dict( + seq_lens=(1,) * n, + slots=ENGINE_SLOTS, + indices=_scattered(n), + ab="separate", + layout="rows", + pool_dtype=torch.bfloat16, + ), + ) + for n in (1, 3, 10, 11, 16, 32) + ) + + ( + Cell( + why="bf16 rows form, pad rows (-1) between slot 0 and 460", + spec=dict( + seq_lens=(1, 1, 1, 1), + slots=ENGINE_SLOTS, + indices=(0, -1, ENGINE_SLOTS - 1, -1), + ab="separate", + layout="rows", + pool_dtype=torch.bfloat16, + ), + ), + Cell( + why="fp32 pool, rows form N=8, 461-slot pool", + spec=dict( + seq_lens=(1,) * 8, + slots=ENGINE_SLOTS, + indices=_scattered(8), + ab="separate", + layout="rows", + ), + ), + ) + tuple( Cell( why=f"decode N={n}, one token per sequence, 24 slots", @@ -171,20 +211,35 @@ class FusedSigmoidGatingDeltaRuleUpdate(OpWrapper): note = """ Metadata consumed: none. - Call site (`_gdn_block` decode): `(a_log, a, dt_bias, 1.0, 20.0, q, k, v, - b, ssm_states, si_all, None, True, qsl_long, None)` -- q/k/v reshaped - column views of the conv output `[T, 10240]` (read through `stride(1)`, no - repack), `a`/`b` two contiguous `[T, 48]` projections, the layer's whole - fp32 pool `[slots, 48, 128, 128]`, int32 slots, int64 `cu_seqlens`, one - token per sequence. Certified at 24 and 461 slots; at 461 output and state - are bitwise identical to 24 for the same slot content. + Call sites (`_gdn_block` decode), one token per sequence, int32 slots, the + layer's whole pool `[slots, 48, 128, 128]` at `mamba_ssm_cache_dtype`: + - fp32 pool: `(a_log, a, dt_bias, 1.0, 20.0, q, k, v, b, ssm_states, + si_all, None, True, qsl_long, None)` -- q/k/v `[1, T, heads, 128]` + reshaped column views of the conv output `[T, 10240]` (read through + `stride(1)`, no repack), `a`/`b` two contiguous `[T, 48]` projections, + int64 `cu_seqlens`; + - bf16 pool: the rows form, `cu_seqlens` None -- q/k/v `[T, 1, heads, + 128]` split from the whole `[T, 10240]` block reshaped to `[T, 1, 80, + 128]`, so the token stride is the row stride (`is_valid` refuses other + strides), `a`/`b` `[T, 1, 48]`. + Certified at 24 and 461 slots; at 461 output and state are bitwise + identical to 24 for the same slot content. The rows form is bitwise the + packed Triton call at either pool dtype and runs inside a CUDA graph. - Branch: the callable takes a FlashInfer CuTe-DSL kernel when the pool is - bf16, `cu_seqlens` is given and `T == N` (head dims 128, supported arch); - everything else -- every fp32-pool call, i.e. this target's - `--mamba_ssm_cache_dtype=float32` -- runs the vendored Triton kernel. Only - the Triton branch is certified; the test proves the fp32 pool never - reaches FlashInfer and that the witness is not blind. + Branch: the callable takes a FlashInfer CuTe-DSL kernel + (`gdn_decode_bf16_state.gated_delta_rule`) when the pool is bf16, + `cu_seqlens` is given and `T == N` (head dims 128, supported arch); + everything else runs the vendored Triton kernel, and only the Triton + branch is certified. The target's two forms both stay on it, which the + test proves with the FlashInfer arm armed to raise. The bf16 call takes + the rows form because FlashInfer is not batch-invariant: at 48 value + heads it runs one kernel up to N = 10 and another from N = 11, and one + fixed row's single-step output and state differ between the two, so a + row's tokens would depend on its batch. The Triton rows are one program + per (row, head); the test shows a fixed row's output and state bitwise + equal across batch sizes 1..32. FlashInfer also differs from `reference` + in reading and writing a negative slot as slot 0 (Triton gives a zero + state and no write) and in its L2-norm epsilon. Preconditions not checked by anything: - `HV % H == 0`. @@ -207,8 +262,8 @@ class FusedSigmoidGatingDeltaRuleUpdate(OpWrapper): not measured: N > 1365 at this geometry). Not certified: the FlashInfer branch, K or V other than 128, fp16/fp32 - operands, `softplus_beta`/`softplus_threshold` other than 1.0/20.0, - `B > 1` with `cu_seqlens`. + operands, an fp16 pool, `softplus_beta`/`softplus_threshold` other than + 1.0/20.0, `B > 1` with `cu_seqlens`. Hard limits none probed. @@ -301,8 +356,8 @@ def reference( ) beta = torch.sigmoid(b3[row, t].float()) if use_qk_l2norm_in_kernel: - q_t = q_t / (q_t.pow(2).sum(-1, keepdim=True).sqrt() + 1e-6) - k_t = k_t / (k_t.pow(2).sum(-1, keepdim=True).sqrt() + 1e-6) + q_t = q_t / torch.sqrt(q_t.pow(2).sum(-1, keepdim=True) + 1e-6) + k_t = k_t / torch.sqrt(k_t.pow(2).sum(-1, keepdim=True) + 1e-6) state = state * torch.exp(neg_exp_a * softplus)[:, None, None] u = (v[row, t].float() - (state * k_t[:, :, None]).sum(1)) * beta[:, None] state = state + k_t[:, :, None] * u[:, None, :] @@ -338,6 +393,16 @@ def is_valid( f"{tuple(output.shape)} -- the kernel writes it with packed " "[NK, B, T, HV, V] arithmetic, misplacing values and clobbering the gaps" ) + if cu_seqlens is None: + # q/k/v reshaped to [N, 1, heads, 128] from column slices of the packed + # [N, 10240] block carry stride(1) = heads * 128, not the row stride: they are + # accepted and read other rows' columns. + for name, t in (("q", q), ("k", k), ("v", v), ("a", a), ("b", b)): + assert t.dim() < 3 or t.stride(0) == t.shape[1] * t.stride(1), ( + f"{name} {tuple(t.shape)} has strides {t.stride()}: without " + "cu_seqlens the kernel finds row n's tokens at n * T * stride(1), " + "so stride(0) must equal T * stride(1)" + ) fused_sigmoid_gating_delta_rule_update = FusedSigmoidGatingDeltaRuleUpdate() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/reset_recurrent_state_rows.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/reset_recurrent_state_rows.py index 5416d7bc8808..7b4467a2cf7c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/reset_recurrent_state_rows.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/reset_recurrent_state_rows.py @@ -22,9 +22,10 @@ #: [conv_dim, d_conv - 1] = [10240, 3] convolution window. REC_SHAPE = (48, 128, 128) CONV_SHAPE = (10240, 3) -#: Largest slot count the cache manager was seen to allocate both pools at for -#: this checkpoint (10 .. 461 over twelve runtime-value dumps, batch never above -#: 32). The target hands both pools over unsliced. +#: A slot count of a full-size engine run, shared by the recurrent and conv pools. The +#: kernel addresses a slot by its index alone, so a pool of another size (a bf16 recurrent +#: pool holds more slots in the same memory) takes the same path. The target hands both +#: pools over unsliced. ENGINE_SLOTS = 461 @@ -80,6 +81,18 @@ class ResetRecurrentStateRows(OpWrapper): why="engine pool 461, 32 requests, slot 0 reset", spec=_engine_case(32, flip=True), ), + Cell( + why="bf16 recurrent pool 461, 4 requests, slot 0 reset", + spec=dict(_engine_case(4, flip=False), rec_dtype=torch.bfloat16), + ), + Cell( + why="bf16 recurrent pool 461, 32 requests, slot 0 reset", + spec=dict(_engine_case(32, flip=False), rec_dtype=torch.bfloat16), + ), + Cell( + why="bf16 recurrent pool 461, 32 requests, flags flipped: slot 460 reset", + spec=dict(_engine_case(32, flip=True), rec_dtype=torch.bfloat16), + ), Cell( why="every request continues: pools bitwise unchanged", spec=dict(slots=6, indices=(0, 1, 2, 3, 4, 5), flags=(True,) * 6), diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/patch_embed_conv3d.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/patch_embed_conv3d.py index 7f9d3b609836..9a6fa3653b37 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/patch_embed_conv3d.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/patch_embed_conv3d.py @@ -60,6 +60,12 @@ class PatchEmbedConv3d(OpWrapper): ) for n in (1, 2, 7, 64, 1024, 4096) ) + + ( + Cell( + why="tower call ceiling: 8192 tokens x merge 2x2 = 32768 patches bf16", + spec=dict(num_patches=32768, dtype=torch.bfloat16), + ), + ) + ( Cell( why="explicit [N, C, kt, kh, kw] rows", @@ -103,7 +109,8 @@ class PatchEmbedConv3d(OpWrapper): note = """ Metadata consumed: none. `pixel_values` is not mutated. All three operands share one dtype and device -- torch raises "... should be the same" - otherwise. Several images are just more rows; there is no batch axis. + otherwise. Several images, and a video's frame groups, are just more rows; + there is no batch axis. fp32 runs in TF32 unless the caller sets `torch.backends.cudnn.allow_tf32 = False` (global state the caller owns): measured 5.84e-4 from an exact fp32 @@ -114,7 +121,12 @@ class PatchEmbedConv3d(OpWrapper): `[1152, 1536]` reshape of the filter is a view of the same bytes. Not timed. CUDA-graph capturable; the output is a fresh allocation. - Not certified: float8 operands; patch counts past 4096. + Not batch-invariant: cuDNN picks its kernel by row count, so a row's last bits + depend on how many rows share the call. Check rows against `reference`, never + against a call of a different size. + + Not certified: float8 operands; patch counts between the 4096- and 32768-row + cells, and past 32768. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/weights.py index e839617d3405..0b02e1bfb7e6 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/weights.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/weights.py @@ -10,8 +10,11 @@ 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); +* NVFP4 routed-expert stacks, laid out for the trunk's routed backend + (`core.trunk_moe_backend`): `[up ; gate]`, then 128x4 scale swizzle. trtllm-gen also + interleaves FC1 row by row and 32-row block shuffles weight and scale bytes (FC2: + shuffle only); CuteDSL interleaves FC1 in 64-row up/gate groups and leaves FC2 in + checkpoint row order; * n-gram table: 128 row chunks concatenated into one `[rows, 160]` buffer (the hash emits global row indices). @@ -23,9 +26,9 @@ more slots; several entries certify at most 461 slots, so raising `kv_cache_fraction` with the offload on needs those cells re-checked. -Not loaded: `mtp.*` when there is no `speculative_config`, and -`mtp.layers.0.self_attn.indexer.*` always (the draft layer runs dense attention). -MTP experts are plain row-major FP8 128x128: `[up ; gate]` concat, bf16 +Not loaded: `mtp.*` when there is no `speculative_config`. With one, the whole MTP module +is loaded, including the draft layer's own QSA indexer (`_index_branch`). MTP experts are +plain row-major FP8 128x128: `[up ; gate]` concat, bf16 `weight_scale_inv` cast to fp32. `lm_head` is the shell's parameter. `load(model, weights)` fills every declared parameter once and asserts coverage @@ -176,15 +179,6 @@ def _host_resident_empty(rows: int, cols: int, dtype, device) -> tuple[torch.Ten _MTP = "mtp" _MTP_LAYER = "mtp.layers.0" -# The three tensors of the draft layer that are never loaded, with their reason -# in the module docstring: its own QSA index projection and the two per-head -# norms that go with it. The draft layer runs dense attention. -_MTP_NOT_LOADED = ( - f"{_MTP_LAYER}.self_attn.indexer.index_qk_proj.weight", - f"{_MTP_LAYER}.self_attn.indexer.q_layernorm.weight", - f"{_MTP_LAYER}.self_attn.indexer.k_layernorm.weight", -) - # NVFP4 block size: one e4m3 scale per 16 contiguous elements along K. SF_VEC = 16 @@ -221,8 +215,25 @@ def _interleave_perm(rows: int, device) -> torch.Tensor: return p.to(device) -def _fc1_perm(rows: int, device) -> torch.Tensor: - """Interleave then 32-row block shuffle, composed into one gather.""" +#: The group the CuteDSL FC1 epilogue pairs up and gate rows in: rows `[128k, 128k + 64)` +#: of the fused operand are up rows and the next 64 their gate rows. +_CUTEDSL_GATE_UP_GROUP = 64 + + +def _group_interleave_perm(rows: int, group: int, device) -> torch.Tensor: + """Gather index that interleaves the `[up | gate]` halves of a `2*I` row stack in + `group`-row blocks: up[0:g], gate[0:g], up[g:2g], gate[g:2g], ...""" + assert rows % (2 * group) == 0, (rows, group) + idx = torch.arange(rows, device=device).view(2, rows // (2 * group), group) + return idx.transpose(0, 1).reshape(-1) + + +def _fc1_perm(rows: int, device, core) -> torch.Tensor: + """The FC1 row gather of the trunk's routed backend. trtllm-gen: interleave then + 32-row block shuffle, composed into one gather. CuteDSL: the 64-row group + interleave alone.""" + if core.trunk_moe_backend == "CUTEDSL": + return _group_interleave_perm(rows, _CUTEDSL_GATE_UP_GROUP, device) return _interleave_perm(rows, device)[_block32_perm(rows, device)] @@ -243,28 +254,34 @@ def _expert_fc1_w(t: tuple[torch.Tensor, ...], core) -> torch.Tensor: """(up, gate) packed nibbles -> one expert's FC1 operand.""" up, gate = t w = torch.cat([up, gate], dim=0) - return torch.index_select(w, 0, _fc1_perm(w.shape[0], w.device)).contiguous() + return torch.index_select(w, 0, _fc1_perm(w.shape[0], w.device, core)).contiguous() def _expert_fc1_s(t: tuple[torch.Tensor, ...], core) -> torch.Tensor: - """(up, gate) e4m3 scales -> one expert's FC1 scale operand.""" + """(up, gate) e4m3 scales -> one expert's FC1 scale operand: the weight's own row + permutation, then the swizzle.""" up, gate = t s = torch.cat([up.view(torch.uint8), gate.view(torch.uint8)], dim=0) rows, cols = s.shape - s = torch.index_select(s, 0, _fc1_perm(rows, s.device)) + s = torch.index_select(s, 0, _fc1_perm(rows, s.device, core)) return _swizzle(s).reshape(rows, cols) def _expert_fc2_w(t: torch.Tensor, core) -> torch.Tensor: - """`down_proj` packed nibbles -> one expert's FC2 operand (no interleave).""" + """`down_proj` packed nibbles -> one expert's FC2 operand (no interleave): the + 32-row block shuffle on trtllm-gen, checkpoint row order on CuteDSL.""" + if core.trunk_moe_backend == "CUTEDSL": + return t.contiguous() return torch.index_select(t, 0, _block32_perm(t.shape[0], t.device)).contiguous() def _expert_fc2_s(t: torch.Tensor, core) -> torch.Tensor: - """`down_proj` e4m3 scales -> one expert's FC2 scale operand.""" + """`down_proj` e4m3 scales -> one expert's FC2 scale operand: the weight's own row + order, then the swizzle.""" s = t.view(torch.uint8) rows, cols = s.shape - s = torch.index_select(s, 0, _block32_perm(rows, s.device)) + if core.trunk_moe_backend != "CUTEDSL": + s = torch.index_select(s, 0, _block32_perm(rows, s.device)) return _swizzle(s).reshape(rows, cols) @@ -473,6 +490,31 @@ def _gated_attention(tag, src, layers): ) +def _index_branch(tag, src, layers): + """The QSA index branch's three weights: one replicated, unquantized projection + producing `[index_heads q | 1 k]`, and two per-head gemma norms.""" + return ( + W( + f"{tag}idx_w", + shape=lambda d: (d.index_width, d.hidden), + src=src("index_qk_proj.weight"), + layers=layers, + ), + W( + f"{tag}idx_qn", + shape=lambda d: (d.index_head_dim,), + src=src("q_layernorm.weight"), + layers=layers, + ), + W( + f"{tag}idx_kn", + shape=lambda d: (d.index_head_dim,), + src=src("k_layernorm.weight"), + layers=layers, + ), + ) + + def _shared_expert(tag, prefix, layers): """The bf16 shared expert beside the routed ones, and its per-token gate.""" return ( @@ -529,26 +571,7 @@ class Qwen4ExpTp1Weights: ), # --- the gated GQA layers and their QSA index branch --- *_gated_attention("", lambda s: f"{{p}}.self_attn.{s}", _attn_layers), - # One replicated, unquantized projection producing [index_heads q | 1 k], and - # two per-head gemma norms. - W( - "idx_w", - shape=lambda d: (d.index_width, d.hidden), - src="{p}.self_attn.indexer.index_qk_proj.weight", - layers=_attn_layers, - ), - W( - "idx_qn", - shape=lambda d: (d.index_head_dim,), - src="{p}.self_attn.indexer.q_layernorm.weight", - layers=_attn_layers, - ), - W( - "idx_kn", - shape=lambda d: (d.index_head_dim,), - src="{p}.self_attn.indexer.k_layernorm.weight", - layers=_attn_layers, - ), + *_index_branch("", lambda s: f"{{p}}.self_attn.indexer.{s}", _attn_layers), # --- the Gated-DeltaNet layers --- W( "in_qkv_w", @@ -872,6 +895,9 @@ class Qwen4ExpTp1Weights: ) ), *_gated_attention("mtp_", lambda s: _mtp_key(f"{_MTP_LAYER}.self_attn.{s}"), _mtp_only), + *_index_branch( + "mtp_", lambda s: _mtp_key(f"{_MTP_LAYER}.self_attn.indexer.{s}"), _mtp_only + ), W( "mtp_router", shape=lambda d: (d.num_experts, d.hidden), @@ -973,13 +999,10 @@ def dims(self, core) -> SimpleNamespace: ) def expected_unconsumed(self, core, weights) -> set: - """With MTP off, the whole `mtp.` subtree; with it on, exactly the three draft - indexer tensors. Named explicitly rather than by prefix so that a fourth unloaded - `mtp.*` key fails the coverage assert instead of hiding.""" + """With MTP off, the whole `mtp.` subtree; with it on, nothing, so that any + unloaded `mtp.*` key fails the coverage assert instead of hiding.""" if core.mtp_enabled: - for key in _MTP_NOT_LOADED: - assert key in weights, f"checkpoint key missing: {key}" - return set(_MTP_NOT_LOADED) + return set() return {k for k in weights if k.startswith(f"{_MTP}.")} def declare(self, d: SimpleNamespace) -> nn.ParameterDict: @@ -1082,9 +1105,7 @@ def fill(param: nn.Parameter, ckpt_key, index, transform) -> None: torch.cuda.synchronize() leftover = set(weights.keys()) - consumed expected = self.expected_unconsumed(core, weights) - what = ( - "the draft layer's three QSA indexer tensors" if core.mtp_enabled else "the MTP module" - ) + what = "empty" if core.mtp_enabled else "the MTP module" assert leftover == expected, ( f"checkpoint coverage: leftover keys are not exactly {what}; " f"unexpected {sorted(leftover - expected)[:8]}, " diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/routing.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/routing.py index 267a7c3b7ff6..2cecbbc5fad0 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/routing.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/routing.py @@ -14,9 +14,12 @@ from typing import Any, Optional, Tuple +import torch + from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig +from tensorrt_llm.quantization.mode import QuantAlgo -from ..._router_index import NoTarget +from ..._router_index import EngineFacts, NoTarget # (layers, hidden, experts, attention heads, linear value heads, vocab). Shape-sniffing # also matches fine-tunes; the accuracy gate pins identity. @@ -27,6 +30,77 @@ # list this target lacks. _VISION_SHAPE = (27, 1152, 16, 2, 2560, 2304, (), "gelu_pytorch_tanh") +# `moe_config.backend` requests the target implements (`_TRUNK_MOE_BACKENDS` in modeling.py). +_MOE_BACKENDS = ("AUTO", "TRTLLM", "CUTEDSL") + +# The longest MTP draft the target runs (`_MAX_VERIFY_WIDTH - 1` in modeling.py); the legacy +# per-step recurrent states (`_MAX_LEGACY_VERIFY_WIDTH - 1`) and the CUTEDSL experts +# (`_CUTEDSL_MAX_DRAFT_LEN`) are certified at shorter ones. +_MAX_DRAFT_LEN = 7 +_MAX_LEGACY_DRAFT_LEN = 3 +_MAX_CUTEDSL_DRAFT_LEN = 3 + +# Recurrent-state cache dtypes the target serves (`_SSM_STATE_DTYPES` in modeling.py). +_SSM_STATE_DTYPES = (torch.float32, torch.bfloat16) + + +def _check_engine( + spec_config: Optional[Any], kv_cache_quant_algo: Optional[Any], engine: Optional[EngineFacts] +) -> None: + """Decline what the target's constructor or first forward would refuse. + + Without this ``auto`` claims the deployment and then fails mid-build instead of + falling back to the built-in implementation. + """ + if spec_config is not None and ( + getattr(spec_config, "draft_len_schedule", None) is not None + or getattr(spec_config, "acceptance_rate_threshold", None) is not None + ): + raise NoTarget( + "the draft length changes at runtime (draft_len_schedule, max_concurrency or " + "acceptance_rate_threshold), but this target's verify step runs at the one " + "width max_draft_len + 1" + ) + if kv_cache_quant_algo not in (None, QuantAlgo.FP8, QuantAlgo.NVFP4): + raise NoTarget( + f"a {kv_cache_quant_algo} KV cache; the target stores K/V as the checkpoint's bf16, " + "as FP8 or as NVFP4" + ) + if engine is None: + return + if engine.torch_dtype != torch.bfloat16: + raise NoTarget(f"the engine runs the model in {engine.torch_dtype}; the target is bf16") + if str(engine.moe_backend).upper() not in _MOE_BACKENDS: + raise NoTarget( + f"moe_config.backend={engine.moe_backend!r}; the target runs the trunk's routed " + f"experts on {', '.join(_MOE_BACKENDS)}" + ) + if engine.mamba_ssm_cache_dtype not in _SSM_STATE_DTYPES: + raise NoTarget( + f"a {engine.mamba_ssm_cache_dtype} recurrent-state cache; the target serves " + "float32 or bfloat16 (kv_cache_config.mamba_ssm_cache_dtype)" + ) + if engine.prefill_cuda_graph_backend == "piecewise": + raise NoTarget( + "prefill_cuda_graph_backend='piecewise' makes the engine torch.compile the model, " + "which the target does not implement; 'breakable' runs the prefill step as a " + "breakable CUDA graph" + ) + if spec_config is None: + return + draft = spec_config.max_draft_len + if not engine.gdn_replay and draft > _MAX_LEGACY_DRAFT_LEN: + raise NoTarget( + f"max_draft_len={draft} with TRTLLM_USE_GDN_REPLAY=0: the legacy per-step recurrent " + f"states are certified at max_draft_len up to {_MAX_LEGACY_DRAFT_LEN}; leave " + "TRTLLM_USE_GDN_REPLAY unset for a longer draft" + ) + if str(engine.moe_backend).upper() == "CUTEDSL" and draft > _MAX_CUTEDSL_DRAFT_LEN: + raise NoTarget( + f"max_draft_len={draft} with moe_config.backend=CUTEDSL: the CuteDSL expert " + f"pipeline is certified at max_draft_len up to {_MAX_CUTEDSL_DRAFT_LEN}" + ) + def route( pretrained_config: Any, @@ -34,6 +108,7 @@ def route( compute_capability: Tuple[int, int], spec_config: Optional[Any], kv_cache_quant_algo: Optional[Any], + engine: Optional[EngineFacts] = None, ) -> type: """Return the target class for this deployment, or raise ``NoTarget`` saying why none fits. @@ -73,17 +148,16 @@ def route( ) # The target implements the one-model MTP-Eagle loop over the checkpoint's - # single MTP module, and no other draft path. A draft longer than three - # needs predicted_tokens_per_seq past the 1..4 thop_attention certifies, - # which the target asserts; declining here falls back instead of crashing. - # The draft must also be a linear chain over that one module; the target - # asserts each of these at construction, so declining here is what lets - # ``auto`` fall back instead of crashing. + # single MTP module, and no other draft path. A draft longer than seven needs + # more verify rows per request than the GDN cached replay and thop_attention + # are certified at. The draft must also be a linear chain over that one + # module; the target asserts each of these at construction, so declining here + # is what lets ``auto`` fall back instead of crashing. mtp_eagle_or_none = spec_config is None or ( isinstance(spec_config, MTPDecodingConfig) and spec_config.spec_dec_mode.is_mtp_eagle_one_model() and getattr(spec_config, "num_nextn_predict_layers", None) in (None, 1) - and 1 <= spec_config.max_draft_len <= 3 + and 1 <= spec_config.max_draft_len <= _MAX_DRAFT_LEN and not getattr(spec_config, "use_dynamic_tree", False) and getattr(spec_config, "eagle_choices", None) is None and spec_config.max_total_draft_tokens == spec_config.max_draft_len @@ -96,8 +170,6 @@ def route( and vshape == _VISION_SHAPE and mapping.world_size == 1 and mtp_eagle_or_none - # The target runs the checkpoint's bf16 KV pool only; it asserts so at construction. - and kv_cache_quant_algo is None ): # One line on purpose: the claims suite reads the routed targets off # `from ..modeling import ` lines. @@ -107,6 +179,7 @@ def route( # isort: on # fmt: on + _check_engine(spec_config, kv_cache_quant_algo, engine) return ModelingV2Qwen38FlashNextNvfp4Sm103Tp1 raise NoTarget( @@ -116,6 +189,5 @@ def route( f"{f' max_draft_len={spec_config.max_draft_len}' if spec_config is not None else ''}, " f"KV cache {kv_cache_quant_algo or 'bf16'}; " "the one qwen4_exp target is Qwen3.8-Flash-Next-NVFP4 on sm103, tp1, with " - "one-model MTP-Eagle (max_draft_len 1..3) or no speculative decoding, and a " - "bf16 KV cache" + f"one-model MTP-Eagle (max_draft_len 1..{_MAX_DRAFT_LEN}) or no speculative decoding" ) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_canonical_selection.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_canonical_selection.py new file mode 100644 index 000000000000..3d25d1ad2efd --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_canonical_selection.py @@ -0,0 +1,270 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qsa_canonical_selection entry. + +The rewrite fixes every slot of every row, so `_compare` is index equality against a +composite-key top-k. The input rows come from an adversary that keeps a random part of the +tie class at the k-th value in a random order, or from the radix engine the target calls. +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_canonical_selection import ( # noqa: E501 + qsa_canonical_selection as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_topk_select import ( # noqa: E501 + qsa_topk_select, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +assert torch.cuda.is_available(), "qsa_canonical_selection requires a CUDA device" + +DEV = "cuda" +COMPRESS_RATIO = 4 +BLOCK_TOPK = 512 +COLS = 2080 +#: What a padding slot sorts under when a raw selection is put in ascending order: past +#: every column, so padding stays last. +PAD_KEY = torch.iinfo(torch.int32).max + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Index for index. The rule fixes every slot of every row, so nothing is tolerated: a + selection does no arithmetic.""" + assert out.dtype == torch.int32, out.dtype + assert out.shape == ref.shape, (out.shape, ref.shape) + wrong = (out != ref).any(1) + if bool(wrong.any()): + row = int(wrong.nonzero()[0, 0]) + raise AssertionError( + f"{int(wrong.sum())} of {out.shape[0]} rows differ from the canonical " + f"selection; first is row {row}" + ) + + +def _logits(kind: str, rows: int, cols: int, gen: torch.Generator) -> torch.Tensor: + if kind == "relu": + return torch.relu(torch.randn(rows, cols, generator=gen, device=DEV)) * 3.0 + if kind == "relu_sparse": + x = torch.relu(torch.randn(rows, cols, generator=gen, device=DEV)) * 3.0 + return x * (torch.rand(rows, cols, generator=gen, device=DEV) < 0.2) + if kind == "signed_zeros": + x = _logits("relu_sparse", rows, cols, gen) + negative = torch.rand(rows, cols, generator=gen, device=DEV) < 0.5 + return torch.where((x == 0) & negative, torch.full_like(x, -0.0), x) + if kind == "two_valued": + return torch.randint(0, 2, (rows, cols), generator=gen, device=DEV).float() + if kind == "all_equal": + return torch.full((rows, cols), 1.5, device=DEV) + if kind == "pair_at_k": + x = torch.rand(rows, cols, generator=gen, device=DEV).argsort(1).float() + 1.0 + order = torch.argsort(x[:, :2048], dim=1, descending=True) + kth, nxt = order[:, BLOCK_TOPK - 1 : BLOCK_TOPK], order[:, BLOCK_TOPK : BLOCK_TOPK + 1] + x.scatter_(1, nxt, torch.gather(x, 1, kth)) + return x + raise ValueError(kind) + + +def _windows(spec: dict, rows: int, cols: int, gen: torch.Generator): + kind = spec["windows"][0] + starts = torch.zeros(rows, dtype=torch.int32, device=DEV) + if kind in ("chunk", "positions_random"): + if kind == "chunk": + positions = spec["windows"][1] + torch.arange(rows, device=DEV) + else: + lo, hi = spec["windows"][1:] + positions = torch.randint(lo, hi + 1, (rows,), generator=gen, device=DEV) + ends = ((positions + 1) // COMPRESS_RATIO).clamp_max(cols).to(torch.int32) + elif kind == "ends": + ends = torch.tensor(spec["windows"][1], dtype=torch.int32, device=DEV) + else: + starts = torch.tensor(spec["windows"][1], dtype=torch.int32, device=DEV) + ends = torch.tensor(spec["windows"][2], dtype=torch.int32, device=DEV) + assert starts.numel() == rows and ends.numel() == rows + return starts, ends + + +def _poison(logits, starts, ends, kind) -> None: + columns = torch.arange(logits.shape[1], device=DEV)[None, :] + outside = (columns < starts[:, None]) | (columns >= ends[:, None]) + if kind == "nan": + logits.masked_fill_(outside, float("nan")) + return + # Values that would win any comparison they were allowed into. + pattern = torch.tensor([float("nan"), float("inf"), 3e38, float("nan")], device=DEV) + logits.copy_(torch.where(outside, pattern[columns % 4].expand_as(logits), logits)) + + +def _order_key(logits: torch.Tensor) -> torch.Tensor: + """int64 that orders as the fp32 values do, -0.0 folded into +0.0.""" + bits = (logits.float() + 0.0).view(torch.int32).long() + return torch.where(bits < 0, bits ^ 0x7FFFFFFF, bits) + + +def _adversary(logits, starts, ends, index_topk, gen) -> torch.Tensor: + """An exact top-k of each window that keeps a random part of the tie class at the + k-th value, in a random order, `-1` last.""" + rows, cols = logits.shape + columns = torch.arange(cols, device=DEV)[None, :] + in_window = (columns >= starts.long()[:, None]) & (columns < ends.long()[:, None]) + tiebreak = torch.rand(rows, cols, generator=gen, device=DEV).argsort(1) + key = torch.where( + in_window, (_order_key(logits) << 32) | tiebreak, torch.iinfo(torch.int64).min + ) + width = min(index_topk, cols) + top = torch.topk(key, width, dim=1).indices + local = torch.where(torch.gather(in_window, 1, top), top - starts.long()[:, None], -1) + shuffle = torch.where(local >= 0, torch.rand(local.shape, generator=gen, device=DEV), 2.0) + local = torch.gather(local, 1, shuffle.argsort(1)) + out = torch.full((rows, index_topk), -1, dtype=torch.int32, device=DEV) + out[:, :width] = local.to(torch.int32) + return out + + +def _radix(logits, starts, ends) -> torch.Tensor: + out = torch.empty(logits.shape[0], BLOCK_TOPK, dtype=torch.int32, device=DEV) + return qsa_topk_select(logits, starts, ends, out, index_topk=BLOCK_TOPK) + + +def _build(spec: dict, rows: int, seed: int): + gen = torch.Generator(device=DEV).manual_seed(seed) + logits = _logits(spec["logits"], rows, spec["cols"], gen) + starts, ends = _windows(spec, rows, spec["cols"], gen) + _poison(logits, starts, ends, spec["tail"]) + if spec["selection"] == "random": + indices = _adversary(logits, starts, ends, spec["index_topk"], gen) + else: + indices = _radix(logits, starts, ends) + return logits, starts, ends, indices + + +def _sorted(raw: torch.Tensor) -> torch.Tensor: + """A raw selection in ascending order, padding last.""" + keyed = torch.where(raw < 0, PAD_KEY, raw) + ordered = torch.sort(keyed, dim=1).values + return torch.where(ordered == PAD_KEY, -1, ordered) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 + certifying(op) + spec = cell.spec + seed = sum(map(ord, cell.why)) + row_counts = spec["rows"] if isinstance(spec["rows"], tuple) else (spec["rows"],) + for rows in row_counts: + args = _build(spec, rows, seed + rows) + expected = op.reference(*args) + out = op(*args) + assert out is args[3], "the rewrite writes `indices` and returns it" + _compare(out, expected) + if not spec.get("graph"): + return + + # Captured once, replayed against new scores, new causal bounds and a new engine + # selection through the same buffers. + logits, starts, ends, indices = _build(spec, rows, seed) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(logits, starts, ends, indices) + for replay in range(2): + fresh = _build(spec, rows, seed + 1000 + replay) + for static, new in zip((logits, starts, ends, indices), fresh): + static.copy_(new) + expected = op.reference(logits, starts, ends, indices) + graph.replay() + torch.cuda.synchronize() + _compare(indices, expected) + + +def test_tie_free_rows_equal_the_radix_engines_sorted_selection() -> None: + """Where the k-th value has no tie, the rewrite is the engine's own set in ascending + order, bit for bit.""" + gen = torch.Generator(device=DEV).manual_seed(4) + rows = 256 + logits = torch.randn(rows, COLS, generator=gen, device=DEV) + positions = torch.randint(0, 8192, (rows,), generator=gen, device=DEV) + ends = ((positions + 1) // COMPRESS_RATIO).clamp_max(COLS).to(torch.int32) + starts = torch.zeros(rows, dtype=torch.int32, device=DEV) + _poison(logits, starts, ends, "nan") + window = torch.where( + torch.arange(COLS, device=DEV)[None, :] < ends[:, None], logits, float("-inf") + ) + top = torch.topk(window, BLOCK_TOPK + 1, dim=1).values + long_rows = ends > BLOCK_TOPK + assert bool((top[long_rows, -2] != top[long_rows, -1]).all()), "a tie at k: reseed" + assert int(long_rows.sum()) > 100 and int((~long_rows).sum()) > 10 + raw = _radix(logits, starts, ends) + expected = _sorted(raw) + assert torch.equal(op(logits, starts, ends, raw), expected) + + +def test_narrowing_the_table_to_a_bound_on_the_windows_changes_nothing() -> None: + """A caller that knows every window ends by column 1500 passes `logits[:, :1500]` of a + wide table: the rows are the full table's, and a 2^24-element table is tiled.""" + gen = torch.Generator(device=DEV).manual_seed(21) + rows, cols, bound = 4096, 8192, 1500 + logits = _logits("relu_sparse", rows, cols, gen) + starts = torch.zeros(rows, dtype=torch.int32, device=DEV) + ends = torch.randint(0, bound + 1, (rows,), generator=gen, device=DEV).to(torch.int32) + _poison(logits, starts, ends, "poison") + raw = _radix(logits, starts, ends) + wide = op(logits, starts, ends, raw.clone()) + narrow = op(logits[:, :bound], starts, ends, raw.clone()) + _compare(wide, op.reference(logits, starts, ends, raw)) + assert torch.equal(wide, narrow) + + +@pytest.mark.parametrize("kind", ["relu_sparse", "two_valued", "eighths", "pair_at_k"]) +def test_radix_launches_differ_and_canonical_rows_do_not(kind) -> None: + """Repeated radix launches over one tie-heavy score matrix. The raw rows differ from + launch to launch -- in order, and where the tie class at the k-th score straddles the + budget, in which tied columns are kept -- and the canonical row is one and the same + row, the reference's, for every launch of every row.""" + gen = torch.Generator(device=DEV).manual_seed(17) + rows, launches = 64, 20 + if kind == "eighths": + # Positive ties: 64 levels over 2048 columns. + logits = torch.randint(0, 64, (rows, COLS), generator=gen, device=DEV).float() / 8 + else: + logits = _logits(kind, rows, COLS, gen) + starts = torch.zeros(rows, dtype=torch.int32, device=DEV) + ends = torch.full((rows,), 2048, dtype=torch.int32, device=DEV) + _poison(logits, starts, ends, "nan") + expected = op.reference(logits, starts, ends, torch.empty(rows, BLOCK_TOPK, device=DEV)) + orders = [set() for _ in range(rows)] + sets = [set() for _ in range(rows)] + canonical = [set() for _ in range(rows)] + for _ in range(launches): + raw = _radix(logits, starts, ends) + for r, (row, ordered) in enumerate(zip(raw.cpu().tolist(), _sorted(raw).cpu().tolist())): + orders[r].add(tuple(row)) + sets[r].add(tuple(ordered)) + out = op(logits, starts, ends, raw) + for r, row in enumerate(out.cpu().tolist()): + canonical[r].add(tuple(row)) + _compare(out, expected) + assert any(len(o) > 1 for o in orders), "radix returned one order on every launch" + assert all(len(c) == 1 for c in canonical) + + +def test_guard_refuses_a_broadcast_row_start() -> None: + """A one-element `row_starts` broadcasts over every row without an error and frames + every row by row 0's start. Shown unguarded first, then refused.""" + gen = torch.Generator(device=DEV).manual_seed(9) + rows = 4 + logits = torch.relu(torch.randn(rows, COLS, generator=gen, device=DEV)) + starts = torch.tensor([0, 10, 20, 30], dtype=torch.int32, device=DEV) + ends = torch.full((rows,), 2000, dtype=torch.int32, device=DEV) + raw = _adversary(logits, starts, ends, BLOCK_TOPK, gen) + expected = op.reference(logits, starts, ends, raw) + out = op(logits, starts[:1], ends, raw.clone()) + assert out.shape == expected.shape and not torch.equal(out, expected), ( + "the broadcast hazard stopped reproducing" + ) + with pytest.raises(AssertionError, match="one entry per logits row"): + with validating(op): + op(logits, starts[:1], ends, raw.clone()) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_expand_block_indices.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_expand_block_indices.py index 3c09c66fc05f..046684168aec 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_expand_block_indices.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_expand_block_indices.py @@ -63,6 +63,12 @@ def _rows(batch) -> tuple[list[int], list[int]]: positions = [4096 + 13 * i for i in range(batch[1])] elif kind == "serving": return _serving_rows(batch[1], batch[2]) + elif kind == "verify": + starts, per_request = batch[1], batch[2] + positions = [s + i for s in starts for i in range(per_request)] + lengths = [s + per_request for s in starts for _ in range(per_request)] + assert max(lengths) <= ENGINE_MAX_SEQ_LEN + return positions, lengths else: positions = list(batch[1]) return positions, [p + 1 for p in positions] @@ -204,9 +210,11 @@ def test_unpacked_prefix_is_silently_dropped() -> None: 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) +@pytest.mark.parametrize("rows", (4, 256)) +def test_cuda_graph_capture_and_replay(rows) -> None: + """Capturable, and a replay recomputes from whatever the buffers now hold: 4 rows, and the + 256 rows a verify step of 32 requests carries at max_draft_len 7.""" + positions = torch.tensor([4096 + 7 * i for i in range(rows)], device=DEV) lengths = (positions + 1).to(torch.int32) blocks = _packed_selection(positions, CR, TOPK, seed=31) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_mtp_reuse_topk.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_mtp_reuse_topk.py new file mode 100644 index 000000000000..538d345d288d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_mtp_reuse_topk.py @@ -0,0 +1,457 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qsa_mtp_reuse_topk catalog entry. + +Captured rows are built the way `qsa_topk_select` emits them: `min(K, visible)` +distinct complete groups in an arbitrary order as a prefix, then `-1`. Inputs +sit in the engine's buffers -- one slot of a `[slots, max_num_sequences, K]` +shared buffer holding stale rows elsewhere, and row slices of the +`max_num_tokens`-tall request map, visible-count and Top-K buffers. Every gate +is bitwise. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_mtp_reuse_topk import ( # noqa: E501 + qsa_mtp_reuse_topk as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +assert torch.cuda.is_available(), "qsa_mtp_reuse_topk requires a CUDA device" + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Bitwise: these are group ids, and one wrong entry attends to a different + group. There is no rounding for a tolerance to absorb.""" + assert out.dtype == ref.dtype, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + torch.testing.assert_close(out, ref, rtol=0, atol=0) + + +DEV = "cuda" +CR = 4 +K = 512 +#: The engine's `max_num_tokens`: height of the request-map and Top-K buffers. +ENGINE_MAX_NUM_TOKENS = 8192 +ENGINE_MAX_SEQ_LEN = 8192 +#: Captured positions for the first requests of a "mixed" batch: empty (0 +#: groups), 1, 100, 255, 511 and 512 groups, rows that gain a group and rows +#: that do not, up to the last position the 8192-token envelope allows. +ANCHORS = (2, 3, 402, 2046, 2050, 8188, 2051, 2047, 1022, 4095, 5998, 7001, 37, 2044, 3001, 6143) + + +def _captured_row(visible: int, block_topk: int, gen: torch.Generator) -> torch.Tensor: + """A legal capture: `min(K, visible)` distinct groups below `visible`, then -1.""" + row = torch.full((block_topk,), -1, dtype=torch.int32) + if visible > 0: + take = min(visible, block_topk) + row[:take] = torch.randperm(visible, generator=gen)[:take].to(torch.int32) + return row + + +def _requests(contexts, drafted, gen) -> list[tuple[int, int]]: + """(captured_visible, current_visible) per request.""" + kind = contexts[0] + if kind == "gaps": + return list(contexts[1]) + if kind == "explicit": + positions = list(contexts[1]) + else: + count = contexts[1] + low, high = {"mixed": (0, 8188), "short": (0, 2046), "full": (2047, 8188)}[kind] + fixed = { + "mixed": ANCHORS, + "short": (2, 2046, 2045, 2044, 2043, 402, 3), + "full": (2047, 2050, 8188, 2051), + }[kind] + positions = list(fixed[:count]) + positions += torch.randint(low, high + 1, (count - len(positions),), generator=gen).tolist() + out = [] + for q, position in enumerate(positions): + d = drafted[q % len(drafted)] + position = min(position, ENGINE_MAX_SEQ_LEN - 1 - d) + out.append(((position + 1) // CR, (position + d + 1) // CR)) + return out + + +def _build(spec, seed: int, requests: int | None = None) -> dict: + """Turn a cell spec into the op's keyword arguments, plus the buffers behind them. + + Returns `dict(kwargs=..., buffers=...)`; `buffers` holds the full tensors the + arguments are views of, so a test can check that nothing outside the view moved. + """ + gen = torch.Generator().manual_seed(seed) + block_topk = spec["block_topk"] + contexts = spec["contexts"] + if requests is not None: + contexts = (contexts[0], requests) + pairs = _requests(contexts, spec["drafted"], gen) + num_requests = len(pairs) + per_request = spec["rows_per_request"] + rows = num_requests * per_request + num_slots, slot = spec["slots"] + shared_rows = max(spec["shared_rows"], num_requests) + meta_dtype = spec["metadata_dtype"] + + # Stale rows everywhere: other layers' slots and requests past this batch. + shared = torch.randint(-1, 3000, (num_slots, shared_rows, block_topk), generator=gen) + shared = shared.to(torch.int32) + shared_vis = torch.randint(0, 3000, (num_slots, shared_rows), generator=gen).to(meta_dtype) + for q, (captured, _) in enumerate(pairs): + shared[slot, q] = _captured_row(captured, block_topk, gen) + shared_vis[slot, q] = captured + + order = list(range(num_requests)) + if spec["request_map"] == "permuted": + order = torch.randperm(num_requests, generator=gen).tolist() + request, visible = [], [] + for q in order: + captured, current = pairs[q] + for j in range(per_request): + request.append(q) + # A later row of the same request sees at least as many groups. + visible.append(current + j // CR) + request = torch.tensor(request, dtype=meta_dtype).reshape(rows) + visible = torch.tensor(visible, dtype=meta_dtype).reshape(rows) + + if spec["layout"] == "engine": + height = max(ENGINE_MAX_NUM_TOKENS, rows + 37) + tall_out = torch.randint(-1, 3000, (height, block_topk), generator=gen).to(torch.int32) + tall_req = torch.randint(0, shared_rows, (height,), generator=gen).to(meta_dtype) + tall_vis = torch.randint(0, 3000, (height,), generator=gen).to(meta_dtype) + tall_req[:rows], tall_vis[:rows] = request, visible + buffers = dict( + shared=shared.to(DEV), + shared_vis=shared_vis.to(DEV), + request=tall_req.to(DEV), + visible=tall_vis.to(DEV), + output=tall_out.to(DEV), + ) + else: + assert spec["layout"] == "dense" + buffers = dict( + shared=shared[slot].clone().to(DEV), + shared_vis=shared_vis[slot].clone().to(DEV), + request=request.to(DEV), + visible=visible.to(DEV), + output=torch.full((rows, block_topk), -7, dtype=torch.int32, device=DEV), + ) + slot = None + kwargs = dict( + shared_indices=buffers["shared"] if slot is None else buffers["shared"][slot], + shared_visible_blocks=buffers["shared_vis"] + if slot is None + else buffers["shared_vis"][slot], + request_indices=buffers["request"][:rows], + visible_blocks=buffers["visible"][:rows], + output=buffers["output"][:rows], + ) + return dict(kwargs=kwargs, buffers=buffers, rows=rows) + + +def _assert_legal_selection(out: torch.Tensor, kwargs: dict) -> None: + """What the producer's contract buys: no duplicate group, none past the horizon.""" + if out.shape[0] == 0: + return + ids = out.long() + ordered = torch.sort(ids, dim=1).values + assert not bool(((ordered[:, 1:] == ordered[:, :-1]) & (ordered[:, 1:] >= 0)).any()), ( + "a group id appears twice in a row" + ) + captured = kwargs["shared_visible_blocks"].long()[kwargs["request_indices"].long()] + horizon = torch.maximum(captured, kwargs["visible_blocks"].long()) + assert bool((ids < horizon[:, None]).all()), "a group id reaches past the row's horizon" + assert bool((ids >= -1).all()) + + +def _check_cell(spec, seed: int, requests: int | None = None) -> torch.Tensor: + built = _build(spec, seed, requests) + kwargs, buffers, rows = built["kwargs"], built["buffers"], built["rows"] + before = {name: tensor.clone() for name, tensor in buffers.items()} + out = op(**kwargs) + torch.cuda.synchronize() + assert out.data_ptr() == kwargs["output"].data_ptr(), "the op returns the buffer it wrote" + assert out.dtype == torch.int32 and tuple(out.shape) == (rows, spec["block_topk"]) + _compare(out, op.reference(**kwargs)) + _assert_legal_selection(out, kwargs) + for name, tensor in buffers.items(): + if name == "output": + assert torch.equal(tensor[rows:], before[name][rows:]), "rows past the view moved" + else: + assert torch.equal(tensor, before[name]), f"{name} was written" + return out + + +def _check_graph(spec, seed: int) -> None: + """At each request count: eager, capture, then replay on new data.""" + for requests in spec["graph_batches"]: + first = _build(spec, seed + requests, requests) + kwargs = first["kwargs"] + eager = op(**kwargs).clone() + _compare(eager, op.reference(**kwargs)) + + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + op(**kwargs) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = op(**kwargs) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(captured, eager), f"{requests} requests: replay != eager" + + # The mixed batch's first requests sit at fixed positions, so the new + # data comes from full rows: every request's captured row changes. + refill = {**spec, "contexts": ("full", requests)} + fresh = _build(refill, seed + 1000 + requests)["kwargs"] + for name, tensor in kwargs.items(): + if name != "output": + tensor.copy_(fresh[name]) + graph.replay() + torch.cuda.synchronize() + expected = op.reference(**kwargs) + assert not torch.equal(expected, eager), f"{requests} requests: new data changes nothing" + _compare(captured, expected) + rerun = op(**{**kwargs, "output": torch.empty_like(kwargs["output"])}) + assert torch.equal(captured, rerun), f"{requests} requests: replay != eager on new data" + del graph + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + seed = zlib.crc32(cell.why.encode()) + if cell.spec["graph_batches"]: + _check_graph(cell.spec, seed) + else: + _check_cell(cell.spec, seed) + + +def _spec_of(prefix: str) -> dict: + (cell,) = [c for c in op.CELLS if c.why.startswith(prefix)] + return cell.spec + + +def test_boundary_rows_take_the_documented_branch() -> None: + """The boundary cell's rows land where the contract says, read off the output.""" + spec = _spec_of("8 requests, one per boundary") + built = _build(spec, seed=5) + kwargs = built["kwargs"] + out = op(**kwargs).cpu() + shared = kwargs["shared_indices"].cpu() + # (captured, current) per request: (0,1) (1,1) (100,101) (511,512) (512,513) + # (2047,2048) (513,513) (512,512) + assert out[0, 0] == 0 and bool((out[0, 1:] == -1).all()), "empty row: group 0 at slot 0" + assert torch.equal(out[1], shared[1]), "no new group: a plain copy" + assert torch.equal(out[2, :100], shared[2, :100]) and out[2, 100] == 100 + assert bool((out[2, 101:] == -1).all()), "short row: appended right after the prefix" + assert torch.equal(out[3, :511], shared[3, :511]) and out[3, 511] == 511, "fills slot 511" + for row, group in ((4, 512), (5, 2047)): + assert torch.equal(out[row, :511], shared[row, 1:]) and out[row, 511] == group, ( + "full row: the new group displaces the oldest id and the rest shift to the front" + ) + assert torch.equal(out[6], shared[6]) and torch.equal(out[7], shared[7]) + + +def test_fresh_two_rows_take_the_documented_branch() -> None: + """Draft steps 4..6 (max_draft_len 7) append up to two groups; read off the output. + + A full row loses its two oldest ids and the survivors shift to the front, a 511-id row + loses its oldest id, and a short row keeps every captured id with the new groups right + after the prefix. At fresh 1 a full row gives up only the id in slot 0. + """ + spec = _spec_of("8 requests at fresh 2") + kwargs = _build(spec, seed=23)["kwargs"] + out = op(**kwargs).cpu() + shared = kwargs["shared_indices"].cpu() + captured = kwargs["shared_visible_blocks"].cpu().tolist() + current = kwargs["visible_blocks"].cpu().tolist() + # (captured, current) per request: (512, 514) (100, 102) (511, 513) (2046, 2048) (1, 3) + # (512, 513) (255, 257) (512, 514) + assert [c - a for a, c in zip(captured, current)] == [2, 2, 2, 2, 2, 1, 2, 2] + for row in (0, 3, 7): + first = captured[row] + assert torch.equal(out[row, : K - 2], shared[row, 2:]), "full row: the oldest two go" + assert out[row, K - 2 :].tolist() == [first, first + 1], "full row: the two new groups last" + assert torch.equal(out[2, : K - 2], shared[2, 1 : K - 1]), "511-id row: the oldest id goes" + assert out[2, K - 2 :].tolist() == [511, 512] + assert torch.equal(out[1, :100], shared[1, :100]) and out[1, 100:102].tolist() == [100, 101] + assert bool((out[1, 102:] == -1).all()), "short row: appended right after the prefix" + assert out[4, :3].tolist() == [0, 1, 2] and bool((out[4, 3:] == -1).all()) + assert torch.equal(out[5, : K - 1], shared[5, 1:]) and out[5, K - 1] == 512, ( + "fresh 1 on a full row: only the id in slot 0 goes" + ) + assert torch.equal(out[6, :255], shared[6, :255]) and out[6, 255:257].tolist() == [255, 256] + assert bool((out[6, 257:] == -1).all()) + + +def test_worked_example_small_geometry() -> None: + """Five rows written out by hand at K = 8, through a permuted request map.""" + shared = torch.tensor( + [ + [5, 2, 7, 0, -1, -1, -1, -1], # 4 of 8 groups, 2 more completed + [3, 1, 6, 4, 0, 2, 7, 5], # full, 1 more: displaces its first id + [-1] * 8, # nothing captured yet + [1, 0, 2, -1, -1, -1, -1, -1], # nothing new: a copy + [4, 1, 0, 3, 2, 6, 5, -1], # 7 ids + 2 new: displaces its first id too + ], + dtype=torch.int32, + device=DEV, + ) + captured = torch.tensor([8, 8, 0, 3, 7], dtype=torch.int32, device=DEV) + request = torch.tensor([1, 0, 4, 2, 3], dtype=torch.int32, device=DEV) + current = torch.tensor([9, 10, 9, 1, 3], dtype=torch.int32, device=DEV) + kwargs = dict( + shared_indices=shared, + shared_visible_blocks=captured, + request_indices=request, + visible_blocks=current, + output=torch.full((5, 8), 99, dtype=torch.int32, device=DEV), + ) + out = op(**kwargs) + _compare(out, op.reference(**kwargs)) + assert out.cpu().tolist() == [ + [1, 6, 4, 0, 2, 7, 5, 8], + [5, 2, 7, 0, 8, 9, -1, -1], + [1, 0, 3, 2, 6, 5, 7, 8], + [0, -1, -1, -1, -1, -1, -1, -1], + [1, 0, 2, -1, -1, -1, -1, -1], + ] + + +def test_unpacked_prefix_overwrites_a_valid_id() -> None: + """The append slot is the count of valid ids, so a -1 hole costs a group, silently. + + Prefix packing is a value property, so no metadata guard can check it. + """ + shared = torch.tensor([[3, -1, 1, -1, -1, -1, -1, -1]], dtype=torch.int32, device=DEV) + kwargs = dict( + shared_indices=shared, + shared_visible_blocks=torch.tensor([4], dtype=torch.int32, device=DEV), + request_indices=torch.tensor([0], dtype=torch.int32, device=DEV), + visible_blocks=torch.tensor([5], dtype=torch.int32, device=DEV), + output=torch.empty((1, 8), dtype=torch.int32, device=DEV), + ) + out = op(**kwargs) + _compare(out, op.reference(**kwargs)) + assert out.cpu().tolist() == [[3, -1, 4, -1, -1, -1, -1, -1]] + + +def _guard_case(name: str): + """Legal arguments, then the one layout `is_valid` refuses. + + Returns (intended, kwargs): `intended` is the reference on the legal + arguments -- what a caller passing the bad layout meant to get. + """ + spec = _spec_of("permuted request map") + kwargs = _build({**spec, "slots": (3, 2)}, seed=13, requests=6)["kwargs"] + intended = op.reference(**kwargs) + rows = kwargs["output"].shape[0] + + def strided(tensor): + wide = torch.zeros((tensor.numel() * 2,), dtype=tensor.dtype, device=DEV) + wide[::2] = tensor.reshape(-1) + return wide[::2] + + if name == "output column slice": + kwargs["output"] = torch.zeros((rows, 2 * K), dtype=torch.int32, device=DEV)[:, :K] + elif name == "shared_indices column slice": + wide = torch.full((kwargs["shared_indices"].shape[0], 2 * K), -1, dtype=torch.int32) + wide = wide.to(DEV) + wide[:, :K] = kwargs["shared_indices"] + kwargs["shared_indices"] = wide[:, :K] + elif name == "shared_indices 3-D, middle extent == block_topk": + # Slot 1 meant; the 3-D buffer passes the width check and reads slot 0. + stack = torch.full((2, 8, 8), -1, dtype=torch.int32, device=DEV) + stack[1, :, :2] = torch.tensor([5, 6], dtype=torch.int32, device=DEV) + intended_kwargs = dict( + shared_indices=stack[1], + shared_visible_blocks=torch.full((8,), 7, dtype=torch.int32, device=DEV), + request_indices=torch.arange(4, dtype=torch.int32, device=DEV), + visible_blocks=torch.full((4,), 8, dtype=torch.int32, device=DEV), + output=torch.empty((4, 8), dtype=torch.int32, device=DEV), + ) + intended = op.reference(**intended_kwargs) + kwargs = {**intended_kwargs, "shared_indices": stack} + elif name == "shared_visible_blocks 2-D [slots, seqs]": + whole = torch.zeros((3, kwargs["shared_visible_blocks"].numel()), dtype=torch.int32) + whole = whole.to(DEV) + whole[2] = kwargs["shared_visible_blocks"] + kwargs["shared_visible_blocks"] = whole + else: + field = name.split()[1] + kwargs[field] = strided(kwargs[field]) + return intended, kwargs + + +GUARDS = ( + ("output column slice", "output must be contiguous"), + ("shared_indices column slice", "shared_indices must be one slot's"), + ("shared_indices 3-D, middle extent == block_topk", "shared_indices must be one slot's"), + ("shared_visible_blocks 2-D [slots, seqs]", "shared_visible_blocks must be one slot's"), + ("strided request_indices", "request_indices must be contiguous"), + ("strided visible_blocks", "visible_blocks must be contiguous"), + ("strided shared_visible_blocks", "shared_visible_blocks must be one slot's"), +) + + +@pytest.mark.parametrize("name, message", GUARDS, ids=[g[0] for g in GUARDS]) +def test_guard_refuses_a_silently_wrong_layout(name: str, message: str) -> None: + """Unguarded the op returns a plausible wrong answer; `is_valid` refuses it.""" + intended, kwargs = _guard_case(name) + out = op(**kwargs) + torch.cuda.synchronize() + assert not torch.equal(out, intended), f"{name}: the op got it right after all" + _, kwargs = _guard_case(name) + with validating(op), pytest.raises(ValueError, match=message): + op(**kwargs) + + +def test_op_rejects_malformed_arguments() -> None: + """The callable's own loud errors: why `is_valid` does not restate them.""" + kwargs = _build(_spec_of("dense standalone"), seed=17)["kwargs"] + cases = [ + ("width does not match", dict(shared_indices=kwargs["shared_indices"][:, : K // 2])), + ( + "one request and bound per query row", + dict(request_indices=kwargs["request_indices"][:3]), + ), + ("one request and bound per query row", dict(visible_blocks=kwargs["visible_blocks"][:3])), + ("requires CUDA tensors", dict(shared_indices=kwargs["shared_indices"].cpu())), + ("requires CUDA tensors", dict(output=kwargs["output"].cpu())), + ] + for message, override in cases: + with pytest.raises(ValueError, match=message): + op(**{**kwargs, **override}) + + +def test_discriminates_plausible_wrong_variants() -> None: + """The bitwise gate separates this contract from the ones a rewrite would produce.""" + spec = _spec_of("8 requests, one per boundary") + kwargs = _build(spec, seed=19)["kwargs"] + out = op(**kwargs).cpu() + rows = kwargs["request_indices"].long().cpu() + shared = kwargs["shared_indices"].cpu()[rows] + captured = kwargs["shared_visible_blocks"].long().cpu()[rows] + current = kwargs["visible_blocks"].long().cpu() + + never_displace = shared.clone() # appends only where there is room + always_last = shared.clone() # always writes the last `fresh` slots + displace_last = shared.clone() # a full row gives up its newest ids, not its oldest + for r in range(out.shape[0]): + fresh = int((current[r] - captured[r]).clamp(0, K)) + filled = int((shared[r] >= 0).sum()) + for j in range(fresh): + if filled + j < K: + never_displace[r, filled + j] = captured[r] + j + always_last[r, K - fresh + j] = captured[r] + j + displace_last[r, min(filled, K - fresh) + j] = captured[r] + j + for wrong in (never_displace, always_last, displace_last, shared): + assert not torch.equal(out, wrong.to(torch.int32)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_scores.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_scores.py index b2ce01b5cbb0..c30b7da0d376 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_scores.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_scores.py @@ -3,8 +3,8 @@ """GPU test for the qsa_paged_index_scores catalog entry. The paged state is real: a `[pages, tokens_per_block, 1, dim]` index cache -sized from the pages the table maps, and an int32 block table whose rows are -request *slots* -- more of them than the batch has requests, assigned by a +(a view of a wider page under the NVFP4 cells) sized from the pages the table maps, and an +int32 block table whose rows are request *slots* -- more of them than the batch has requests, assigned by a seeded permutation -- with shuffled physical pages and `-1` past each sequence on the serving layouts, which is what the engine's table carries. @@ -186,9 +186,20 @@ def _build(spec, seed): for slot, count in zip(slot_of.tolist(), mapped): block_table[slot, :count] = page_ids[cursor : cursor + count] cursor += count - index_cache = torch.randn( - pool_pages, spec["tokens_per_block"], 1, head_dim, generator=gen, device=DEV - ).to(dtype) + tpb = spec["tokens_per_block"] + keys = torch.randn(pool_pages, tpb, 1, head_dim, generator=gen, device=DEV).to(dtype) + if "page_pad" in spec: + # The NVFP4 pool's page: the other index buffers, this layer's, then the position + # page, so the plane is a view at page stride buffers * tpb * dim + page_pad. + buffers, buffer = spec["pool_buffers"], tpb * head_dim + page = torch.zeros(pool_pages, buffers * buffer + spec["page_pad"], dtype=dtype, device=DEV) + index_cache = page[:, (buffers - 1) * buffer : buffers * buffer].view( + pool_pages, tpb, 1, head_dim + ) + index_cache.copy_(keys) + assert index_cache.stride() == (page.stride(0), head_dim, head_dim, 1) + else: + index_cache = keys if spec["q_layout"] == "scratch_slice": # [T, index heads | 1 k | 1 v, D], q is the first `heads` slots. diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_store.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_store.py index a8c3b5b8fd53..32b4c1aae597 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_store.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_store.py @@ -53,6 +53,28 @@ def _bits(tensor: torch.Tensor) -> torch.Tensor: return tensor.contiguous().view(as_int) +def _shared_page_planes(spec, extras: dict): + """The NVFP4-KV pool's side planes: each physical page holds the `layers` + index-K pages and then the position page, so the index cache (the last + role) and the position cache are strided views of one byte allocation. + Every role holds `KEY_FILL` and the position page `POS_FILL`; the other + roles are recorded for the untouched check.""" + layers, tpb, hkv, head_dim, axes, pages = ( + spec[k] for k in ("layers", "tokens_per_block", "hkv", "head_dim", "axes", "pages") + ) + index_bytes = tpb * hkv * head_dim * 2 + page = torch.empty(pages, layers * index_bytes + tpb * axes * 4, dtype=torch.uint8, device=DEV) + page[:, : layers * index_bytes].view(torch.bfloat16).fill_(KEY_FILL) + page[:, layers * index_bytes :].view(torch.int32).fill_(POS_FILL) + index_cache = page[:, (layers - 1) * index_bytes : layers * index_bytes].view(torch.bfloat16) + index_cache = index_cache.view(pages, tpb, hkv, head_dim) + position_cache = page[:, layers * index_bytes :].view(torch.int32).view(pages, tpb, axes) + assert index_cache.stride() == (spec["page_stride"], hkv * head_dim, head_dim, 1) + assert position_cache.stride() == spec["position_strides"] + extras["roles"] = page[:, : (layers - 1) * index_bytes] + return index_cache, position_cache + + def _build(spec, seed: int): """Real planes, table and rows for one cell; returns the op's kwargs and extras.""" gen = torch.Generator(device="cpu").manual_seed(seed) @@ -112,7 +134,10 @@ def _build(spec, seed: int): coords = coords.contiguous() extras = {} - if spec["cache_layout"] == "role_view": + position_cache = None + if spec["cache_layout"] == "shared_page": + index_cache, position_cache = _shared_page_planes(spec, extras) + elif 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()) @@ -121,7 +146,8 @@ def _build(spec, seed: int): index_cache = torch.full((pages, tpb, hkv, head_dim), KEY_FILL, device=DEV).to( torch.bfloat16 ) - position_cache = torch.full((pages, tpb, axes), POS_FILL, dtype=torch.int32, device=DEV) + if position_cache is None: + position_cache = torch.full((pages, tpb, axes), POS_FILL, dtype=torch.int32, device=DEV) args = dict( token_k=token_k, position_coordinates=coords, @@ -147,6 +173,11 @@ def test_certified_cells(cell, certifying): # noqa: F811 if "neighbour" in extras: view, saved = extras["neighbour"] assert torch.equal(_bits(view), _bits(saved)), "the neighbouring role moved" + if "roles" in extras: + fill = int(_bits(torch.tensor([KEY_FILL], dtype=torch.bfloat16))[0]) + assert bool((extras["roles"].view(torch.int16) == fill).all()), ( + "another index-K role of the shared page moved" + ) if cell.spec["axes"] > 3: assert torch.equal(out[1][..., 3:], before[..., 3:]), "axes past the coordinates moved" if cell.spec["coords_layout"] == "stride0": @@ -247,9 +278,15 @@ def test_repeats_are_stream_ordered_and_idempotent() -> None: assert torch.equal(args["position_cache"], first[1]) -def test_cuda_graph_capture_and_replay() -> None: +#: The cell of 32 MTP7 verify blocks (256 rows): the segments, pool and table it runs over. +_MTP7_32 = next(c.spec for c in op.CELLS if c.why.startswith("decode MTP7: 32 requests")) +_MTP7_BATCH = {key: _MTP7_32[key] for key in ("segments", "pages", "slots")} + + +@pytest.mark.parametrize("batch", ({}, _MTP7_BATCH), ids=("decode", "mtp7-32x8")) +def test_cuda_graph_capture_and_replay(batch) -> None: """Captured once, replayed against new positions written into the same buffer.""" - args = _small(seed=52) + args = _small(seed=52, **batch) positions = args["logical_positions"] replayed = positions.clone() positions.zero_() diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store.py index 2a1fe4cc9440..a77ccf3190ef 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store.py @@ -101,6 +101,13 @@ def _build(spec, seed: int): ).to(device=DEV, dtype=spec["row_dtype"]) k = qkv[:, spec["heads_q"] : spec["heads_q"] + hkv] v = qkv[:, spec["heads_q"] + hkv :] + if spec["values"] == "every_bf16": + patterns = torch.arange(-(2**15), 2**15, dtype=torch.int32).to(torch.int16) + patterns = patterns.view(torch.bfloat16).reshape(rows, hkv, head_dim).to(DEV) + k.copy_(patterns) + v.copy_(patterns.flip(0)) + else: + assert spec["values"] == "randn" if spec["k_layout"] == "head_major": k = k.transpose(1, 2).contiguous().transpose(1, 2) v = v.transpose(1, 2).contiguous().transpose(1, 2) @@ -162,9 +169,15 @@ def test_fp8_pool_saturates_an_out_of_range_row() -> None: assert float(stored[0, 0].float()) == 448.0 -def test_cuda_graph_replay_follows_updated_positions() -> None: +#: The cell of 32 MTP7 verify blocks (256 rows): the segments, pool and table it runs over. +_MTP7_32 = next(c.spec for c in op.CELLS if c.why.startswith("decode MTP7: 32 requests")) +_MTP7_BATCH = {key: _MTP7_32[key] for key in ("segments", "pages", "slots")} + + +@pytest.mark.parametrize("batch", ({}, _MTP7_BATCH), ids=("decode", "mtp7-32x8")) +def test_cuda_graph_replay_follows_updated_positions(batch) -> None: """Captured once, replayed against new positions written into the same buffer.""" - args, pool = _decode_args(seed=9) + args, pool = _decode_args(seed=9, **batch) positions = args["logical_positions"] saved = positions.clone() positions.zero_() diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store_nvfp4.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store_nvfp4.py new file mode 100644 index 000000000000..6c934acfed42 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store_nvfp4.py @@ -0,0 +1,490 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qsa_paged_kv_store_nvfp4 entry. + +The paged state is real: one NVFP4 pool as the QSA cache manager serves it -- e2m1 +data `[pages, 2, Hkv, tpb, D/2]` and e4m3 block scales `[pages, 2, Hkv, tpb, D/16]`, +K and V the two role planes of each -- and the engine's `[slots, width]` table, a +row per request slot, more slots than requests, live rows' pages drawn from a +seeded permutation, `-1` past each allocation. Both pools are pre-filled with a +sentinel so a stray write shows. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_kv_store_nvfp4 import ( # noqa: E501 + qsa_paged_kv_store_nvfp4 as op, +) +from tensorrt_llm.bindings.internal import thop + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +assert torch.cuda.is_available(), "qsa_paged_kv_store_nvfp4 requires a CUDA device" + +DEV = "cuda" +DATA_FILL, SCALE_FILL = 0xA5, 0x5A +QUANT_MODE_NVFP4_KV = 8192 + + +def _compare(out, ref) -> None: + """Bitwise over all four planes: the store is a pure function of the row's + bits, and `reference` performs the same rounding steps.""" + names = ("k_cache", "v_cache", "k_scale_cache", "v_scale_cache") + for name, got, want in zip(names, out, ref): + assert got.dtype == want.dtype and got.shape == want.shape, ( + name, + got.shape, + want.shape, + ) + differ = (got != want).nonzero() + if differ.numel(): + first = tuple(differ[0].tolist()) + raise AssertionError( + f"{name}: {differ.shape[0]} bytes differ from the reference, first at " + f"{first}: {int(got[first]):#04x} vs {int(want[first]):#04x}" + ) + + +def _hole(spec, table, 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 _fill_rows(spec, k, v, gen): + """Row values: random with per-block magnitudes over 2^-10..2^5, every finite bf16, + or edge blocks (zero, -0.0, an outlier).""" + hkv, head_dim = spec["hkv"], spec["head_dim"] + if spec["values"] == "finite_bf16": + bits = torch.arange(-(2**15), 2**15, dtype=torch.int32).to(torch.int16) + finite = bits.view(torch.bfloat16) + finite = finite[torch.isfinite(finite.float())] # 65280 values + padded = torch.zeros(k.numel(), dtype=torch.bfloat16) + assert finite.numel() <= padded.numel() + padded[: finite.numel()] = finite + k.copy_(padded.reshape(k.shape).to(DEV)) + order = torch.randperm(padded.numel(), generator=gen) + v.copy_(padded[order].reshape(v.shape).to(DEV)) + return + rows = k.shape[0] + for t in (k, v): + base = torch.randn(rows, hkv, head_dim // 16, 16, generator=gen) + exponent = torch.randint(-10, 6, (rows, hkv, head_dim // 16, 1), generator=gen) + t.copy_((base * torch.pow(2.0, exponent.float())).reshape(t.shape).to(DEV)) + if spec["values"] == "edge": + k[0, 0, :16] = 0.0 + v[0, 1, 16:32] = -0.0 + k[1, 1, 32:48] = torch.tensor([-0.0] * 8 + [0.0] * 8) + k[2, 0, 5] = 900.0 + v[3, 0, 200] = -900.0 + k[4, 1, 64:80] = 1e-30 # far below the smallest e4m3 scale + else: + assert spec["values"] == "randn" + + +def _pool(spec): + """The allocation the cell's planes view: data and scales of `slot_pages` role + pages per slot, every byte its sentinel.""" + pages, slot_pages = spec["pages"], spec["slot_pages"] + hkv, tpb, head_dim = spec["hkv"], spec["tokens_per_block"], spec["head_dim"] + data = torch.full( + (pages, slot_pages, hkv, tpb, head_dim // 2), DATA_FILL, dtype=torch.uint8, device=DEV + ) + scales = torch.full( + (pages, slot_pages, hkv, tpb, head_dim // 16), SCALE_FILL, dtype=torch.uint8, device=DEV + ) + return data, scales + + +def _planes(args): + return (args["k_cache"], args["v_cache"], args["k_scale_cache"], args["v_scale_cache"]) + + +def _reset(planes) -> None: + for plane, fill in zip(planes, (DATA_FILL, DATA_FILL, SCALE_FILL, SCALE_FILL)): + plane.fill_(fill) + + +def _assert_other_roles_untouched(spec, pool) -> None: + """Every role page of a slot but the call's K and V still holds its sentinel.""" + data, scales = pool + for role in range(spec["slot_pages"]): + if role // 2 != spec["layer"]: + assert bool((data[:, role] == DATA_FILL).all()), f"data role {role} moved" + assert bool((scales[:, role] == SCALE_FILL).all()), f"scale role {role} moved" + + +def _build(spec, seed: int, pool=None): + """Real pool, table and rows for one cell; returns the op's kwargs, the pool, + and the packed QKV rows the K/V slices view. `pool` reuses an allocation + from `_pool(spec)`.""" + gen = torch.Generator(device="cpu").manual_seed(seed) + tpb, hkv, head_dim = spec["tokens_per_block"], spec["hkv"], spec["head_dim"] + slots, width, pages = spec["slots"], spec["width"], spec["pages"] + segments = spec["segments"] + requests = [r for r, (_, n) in enumerate(segments) for _ in range(n)] + positions = [p for start, n in segments for p in range(start, start + n)] + + table = torch.randint(0, pages, (slots, width), generator=gen, dtype=torch.int32) + order = torch.randperm(pages, generator=gen).to(torch.int32) + used, allocated_cols = 0, [] + for r, (start, n) in enumerate(segments): + cols = (start + n - 1) // tpb + 1 + table[r] = -1 + table[r, :cols] = order[used : used + cols] + used += cols + allocated_cols.append(cols) + assert used <= pages and len(segments) <= slots, "pool or table too small for the batch" + if spec["hole"] is not None: + request, position = _hole(spec, table, len(segments), allocated_cols) + requests.append(request) + positions.append(position) + + rows = len(positions) + qkv = torch.zeros( + rows, spec["heads_q"] + 2 * hkv, head_dim, dtype=spec["row_dtype"], device=DEV + ) + k = qkv[:, spec["heads_q"] : spec["heads_q"] + hkv] + v = qkv[:, spec["heads_q"] + hkv :] + qkv[:, : spec["heads_q"]] = torch.randn(rows, spec["heads_q"], head_dim, generator=gen).to( + device=DEV, dtype=spec["row_dtype"] + ) + _fill_rows(spec, k, v, gen) + + layer = spec["layer"] + data, scales = _pool(spec) if pool is None else pool + args = dict( + k=k, + v=v, + k_cache=data[:, 2 * layer], + v_cache=data[:, 2 * layer + 1], + k_scale_cache=scales[:, 2 * layer], + v_scale_cache=scales[:, 2 * layer + 1], + request_indices=torch.tensor(requests, dtype=spec["request_dtype"], device=DEV), + logical_positions=torch.tensor(positions, dtype=spec["position_dtype"], device=DEV), + block_table=table.to(DEV), + tokens_per_block=tpb, + kv_scale_orig_quant=torch.tensor( + [1.0, spec["global_k"], spec["global_v"]], dtype=torch.float32, device=DEV + ), + ) + return args, (data, scales), qkv + + +def _snapshot(pool): + return tuple(t.clone() for t in pool) + + +def _pools_equal(a, b) -> bool: + return all(torch.equal(x, y) for x, y in zip(a, b)) + + +def _graph_replay(spec, seed: int) -> None: + """Each G: capture one step of the first G requests, replay the next step. + + The next step lives in the captured buffers: new k/v rows, new pages in the + table, and every request's positions advanced by its rows, the way a decode + step moves. The replay must equal the reference and an eager call on the + same inputs, bit for bit. One pool serves the whole sweep, as the engine's does. + """ + pool = _pool(spec) + steps = {n for _, n in spec["segments"]} + assert len(steps) == 1, "a graph cell's requests all carry the same row count" + step = steps.pop() + for g in spec["graph_batches"]: + sub = {**spec, "segments": spec["segments"][:g]} + args, _, _ = _build(sub, seed=seed + g, pool=pool) + following, _, _ = _build(sub, seed=seed + 1000 + g, pool=pool) + planes = _planes(args) + positions = args["logical_positions"] + next_positions = positions.clone() + positions.sub_(step) + assert int(positions.min()) >= 0 + + def call(): + op(**args) + + for _ in range(3): # compile the Triton kernel before capture + call() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + call() + + for name in ("k", "v", "block_table"): + args[name].copy_(following[name]) + positions.copy_(next_positions) + _reset(planes) + op(**args) + eager = tuple(plane.clone() for plane in planes) + _reset(planes) + ref = op.reference(**args) + graph.replay() + torch.cuda.synchronize() + _compare(planes, ref) + assert all(torch.equal(p, e) for p, e in zip(planes, eager)), ( + f"G={g} x {step} rows: replay differs from eager" + ) + del graph + _assert_other_roles_untouched(spec, pool) + + +def _thop_append(qkv, pool, table, requests, new, cached, num_contexts, p, scale_orig_quant): + """The attention op's own append over the same rows (e4m3 output: the append + does not depend on the output type).""" + data, scales = pool + tpb = 32 + hq, hkv, d = 24, 2, 256 + ns = len(requests) + width = table.shape[1] + offsets = torch.zeros(1, ns, 2, width, dtype=torch.int32, device=DEV) + for row, r in enumerate(requests): + live = table[r].clamp_min(0) + offsets[0, row, 0] = 2 * live + offsets[0, row, 1] = 2 * live + 1 + pointers = torch.zeros(1, 2, 2, dtype=torch.int64) + pointers[0, 0, 0], pointers[0, 0, 1] = data.data_ptr(), scales.data_ptr() + kv = [c + n for c, n in zip(cached, new)] + ctx = [n if i < num_contexts else kv[i] for i, n in enumerate(new)] + scale_quant_orig = scale_orig_quant.reciprocal() + rows = qkv.reshape(qkv.shape[0], -1) + thop.attention( + q=rows, k=None, v=None, + output=torch.empty(rows.shape[0], hq * d, dtype=torch.float8_e4m3fn, device=DEV), + output_sf=None, workspace_=torch.empty(0, dtype=torch.int8, device=DEV), + sequence_length=torch.tensor(kv, dtype=torch.int32, device=DEV), + host_past_key_value_lengths=torch.tensor(kv, dtype=torch.int32), + host_total_kv_lens=torch.tensor( + [sum(kv[:num_contexts]), sum(kv[num_contexts:])], dtype=torch.int32 + ), + context_lengths=torch.tensor(ctx, dtype=torch.int32, device=DEV), + host_context_lengths=torch.tensor(ctx, dtype=torch.int32), + host_request_types=torch.tensor( + [0 if i < num_contexts else 1 for i in range(ns)], dtype=torch.int32 + ), + max_context_q_len_override=None, kv_cache_block_offsets=offsets, + host_kv_cache_pool_pointers=pointers, + host_kv_cache_pool_mapping=torch.zeros(1, 2, dtype=torch.int32), + cache_indirection=None, kv_scale_orig_quant=scale_orig_quant, kv_scale_quant_orig=scale_quant_orig, + out_scale=torch.ones(1, dtype=torch.float32, device=DEV), rotary_inv_freq=None, + rotary_cos_sin=None, latent_cache=None, q_pe=None, block_ids_per_seq=None, + attention_sinks=None, is_fused_qkv=True, update_kv_cache=True, + predicted_tokens_per_seq=p, local_layer_idx=0, num_heads=hq, num_kv_heads=hkv, + head_size=d, tokens_per_block=tpb, max_num_requests=ns, + max_context_length=width * tpb, max_seq_len=width * tpb, + attention_window_size=width * tpb, beam_width=1, mask_type=1, + quant_mode=QUANT_MODE_NVFP4_KV, q_scaling=1.0, position_embedding_type=0, rope_dim=0, + rope_base=10000.0, rope_scale_type=0, rope_scale=1.0, rope_short_m_scale=1.0, + rope_long_m_scale=1.0, rope_max_positions=1024, rope_original_max_positions=1024, + use_paged_context_fmha=True, attention_input_type=0, is_mla_enable=False, + chunked_prefill_buffer_batch_size=1, q_lora_rank=None, kv_lora_rank=None, + qk_nope_head_dim=None, qk_rope_head_dim=None, v_head_dim=None, rope_append=None, + mrope_rotary_cos_sin=None, mrope_position_deltas=None, helix_position_offsets=None, + helix_is_inactive_rank=None, attention_chunk_size=None, softmax_stats_tensor=None, + is_spec_decoding_enabled=False, use_spec_decoding=False, is_spec_dec_tree=False, + spec_decoding_generation_lengths=None, spec_decoding_position_offsets_for_cpp=None, + spec_decoding_packed_mask=None, spec_decoding_bl_tree_mask_offset=None, + spec_decoding_bl_tree_mask=None, spec_bl_tree_first_sparse_mask_offset_kv=None, + sparse_kv_indices=None, sparse_kv_offsets=None, sparse_attn_indices=None, + sparse_attn_offsets=None, sparse_attn_indices_block_size=0, num_contexts=num_contexts, + num_ctx_tokens=sum(new[:num_contexts]), + ) # fmt: skip + torch.cuda.synchronize() + + +def _thop_equality(spec, seed: int) -> None: + """Steps of a context, a chunked continuation, decode P=1 and MTP verify P=3 over + three requests, each written once by the attention op's append into one pool and + once by this op into another, at the cell's second-level scales: the two pools stay + bitwise equal, and this op's output equals its reference at every step.""" + base, ours, _ = _build(spec, seed=seed) + theirs = _snapshot(ours) + gen = torch.Generator(device="cpu").manual_seed(seed + 1) + table = base["block_table"] + cached = [0, 0, 0] + for new, nc, p in ( + ([70, 33, 5], 3, 1), + ([41, 64, 27], 3, 1), + ([1, 1, 1], 0, 1), + ([3] * 3, 0, 3), + ): + rows = sum(new) + qkv = torch.randn(rows, 28, 256, generator=gen) + qkv *= torch.pow(2.0, torch.randint(-6, 4, (rows, 28, 1), generator=gen).float()) + qkv = qkv.to(device=DEV, dtype=torch.bfloat16) + requests = [r for r, n in enumerate(new) for _ in range(n)] + positions = [c + i for c, n in zip(cached, new) for i in range(n)] + _thop_append(qkv, theirs, table, [0, 1, 2], new, cached, nc, p, base["kv_scale_orig_quant"]) + args = { + **base, + "k": qkv[:, 24:26], + "v": qkv[:, 26:28], + "request_indices": torch.tensor(requests, dtype=torch.int32, device=DEV), + "logical_positions": torch.tensor(positions, dtype=torch.int64, device=DEV), + } + ref = op.reference(**args) + out = op(**args) + _compare(out, ref) + for plane, (a, b) in enumerate(zip(ours, theirs)): + assert torch.equal(a, b), ( + f"step {new}: {('data', 'scales')[plane]} differ from the attention op's " + f"append in {int((a != b).sum())} bytes" + ) + cached = [c + n for c, n in zip(cached, new)] + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + seed = zlib.crc32(cell.why.encode()) + if cell.spec["graph_batches"]: + _graph_replay(cell.spec, seed=seed) + return + if cell.spec["thop_append"]: + _thop_equality(cell.spec, seed=seed) + return + args, pool, _ = _build(cell.spec, seed=seed) + before = _snapshot(pool) + ref = op.reference(**args) + out = op(**args) + assert out[0] is args["k_cache"] and out[3] is args["v_scale_cache"] + _compare(out, ref) + wrote = not _pools_equal(pool, before) + assert wrote == bool(cell.spec["segments"]), ( + "a cell with rows must write, an empty one must not" + ) + _assert_other_roles_untouched(cell.spec, pool) + + +def _specialization_class(value: int) -> str: + """Triton's class for a plain-int kernel argument: 1, a multiple of 16, or other.""" + if value == 1: + return "one" + return "multiple_of_16" if value % 16 == 0 else "other" + + +def test_cells_reach_every_served_specialization_class() -> None: + """The pool's slot count and the table's height each select a compiled variant + by class. Slot count other or a multiple of 16 against table height 1, other + or a multiple of 16 all run, eagerly.""" + reached = { + (_specialization_class(c.spec["pages"]), _specialization_class(c.spec["slots"])) + for c in op.CELLS + if not c.spec["graph_batches"] + } + wanted = { + (pages, height) + for pages in ("other", "multiple_of_16") + for height in ("one", "other", "multiple_of_16") + } + assert wanted <= reached, f"unreached: {sorted(wanted - reached)}" + + +def _decode_args(seed: int, **overrides): + return _build({**op.CELLS[0].spec, **overrides}, seed=seed) + + +def test_rounding_rivals_are_rejected() -> None: + """The bitwise gate separates the engine's rounding from near misses on the + every-finite-bf16 rows: e2m1 ties rounded away from zero instead of to even, + and V's scales laid out token-major like K's, each move bytes the gate sees.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention import ( # noqa: E501 + qsa_paged_kv_store_nvfp4 as entry, + ) + + spec = next(c.spec for c in op.CELLS if c.spec["values"] == "finite_bf16") + args, _, _ = _build(spec, seed=21) + out = op(**args) + _compare(out, op.reference(**args)) + + ties_up = tuple((threshold, True) for threshold, _ in entry._E2M1_STEPS) + saved = entry._E2M1_STEPS + try: + entry._E2M1_STEPS = ties_up + rival = op.reference(**args) + finally: + entry._E2M1_STEPS = saved + moved = int((rival[0] != out[0]).sum()) + assert moved > 0, "ties away from zero must move some K bytes on every-finite-bf16 rows" + + # V's scale layout, on random rows (the every-finite rows saturate V's scales). + spec = next(c.spec for c in op.CELLS if c.why.startswith("mixed:")) + args, _, _ = _build(spec, seed=22) + out = op(**args) + page = int(args["block_table"][0, 0]) + groups = spec["head_dim"] // 16 + interleaved = out[3][page].reshape(spec["hkv"], -1) + token_major = ( + interleaved.reshape(spec["hkv"], spec["tokens_per_block"] // 4, groups, 4) + .permute(0, 1, 3, 2) + .reshape(spec["hkv"], -1) + ) + assert not torch.equal(interleaved, token_major), "the V scale layouts must be told apart" + + +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): + call() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + call() + pool[0].fill_(DATA_FILL) + pool[1].fill_(SCALE_FILL) + positions.copy_(saved) + ref = op.reference(**args) + graph.replay() + torch.cuda.synchronize() + _compare((args["k_cache"], args["v_cache"], args["k_scale_cache"], args["v_scale_cache"]), ref) + + +def test_op_rejects_operands_it_cannot_serve() -> None: + """The callable's own loud errors: why the entry carries no `is_valid`.""" + args, _, _ = _decode_args(seed=5) + cache, scales = args["k_cache"], args["k_scale_cache"] + cases = { + "k and v must both be BF16 or both FP16": dict(k=args["k"].float(), v=args["v"].float()), + "k_cache must be uint8": dict(k_cache=cache.view(torch.int8)), + "k_scale_cache must be uint8": dict(k_scale_cache=scales[..., :8]), + "tokens_per_block must be a positive multiple of 4": dict(tokens_per_block=30), + "block_table must be a 2-D int32 tensor": dict(block_table=args["block_table"].long()), + "request_indices must be contiguous": dict( + request_indices=torch.stack([args["request_indices"]] * 2, 1)[:, 0] + ), + "kv_scale_orig_quant must be a contiguous fp32": dict( + kv_scale_orig_quant=args["kv_scale_orig_quant"][:1] + ), + "every operand must be on k's CUDA device": dict(block_table=args["block_table"].cpu()), + } + for message, override in cases.items(): + with pytest.raises(ValueError, match=message): + with validating(op): + op(**{**args, **override}) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa.py index f7d69fa1df94..4f94114d1d9c 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa.py @@ -55,6 +55,9 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: DEV = "cuda" ENGINE_MAX_SEQ_LEN = 8192 +FP8 = torch.float8_e4m3fn +#: The e4m3fn NaN bit pattern (and, with the sign bit, 0xFF). +E4M3_NAN = 0x7F _KERNELS = { "_qsa_paged_sparse_gqa_kernel": "prefill", @@ -488,3 +491,84 @@ def test_selection_width_is_checked_on_the_prefill_kernel_only() -> None: 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)) + + +# -------------------------------------------------------------------------- +# an fp8 pool is read at face value +# -------------------------------------------------------------------------- + + +def _every_finite_e4m3_code(): + """The 256 e4m3fn codes as bytes, the two NaN codes replaced by +0.""" + codes = torch.arange(256, dtype=torch.int32, device=DEV).to(torch.uint8) + return torch.where((codes & 0x7F) == E4M3_NAN, torch.zeros_like(codes), codes) + + +def _two_token_rows(kernel, seed): + """Rows at position 1 of a request each, on an fp8 pool: every row's + selection is tokens 0 and 1, both on the request's first page.""" + rows = 300 if kernel == "prefill" else 12 + spec = _with( + segments=((1, 1),) * rows, + slots=rows + 8, + prefill_pair=kernel == "prefill", + cache_dtype=FP8, + ) + args, _ = _build(spec, seed) + pages = args["block_table"][args["request_indices"].long(), 0].long() + return args, pages + + +@pytest.mark.parametrize("kernel", ["prefill", "splitk"]) +def test_fp8_values_are_read_at_face_value(kernel, launches) -> None: + """Token 0's V row holds every finite code, token 1's is zero, and q is + zero, so each output row is exactly half of token 0's row: no scale, and + the widening to bf16 is exact for every code, subnormals included.""" + args, pages = _two_token_rows(kernel, seed=12000) + codes = _every_finite_e4m3_code() + values = args["v_cache"].view(torch.uint8) + values[pages, :, 0, :] = codes + values[pages, :, 1, :] = 0 + args["q"].zero_() + out = op(**args) + assert launches[0][0] == kernel + expected = codes.view(FP8).to(torch.bfloat16) / 2 + assert torch.equal(out, expected.expand_as(out)) + _compare(out, op.reference(**args)) + + +@pytest.mark.parametrize("kernel", ["prefill", "splitk"]) +def test_fp8_keys_are_read_at_face_value(kernel, launches) -> None: + """Token 0's K row holds every finite code, token 1's is zero, and each + (row, head) queries one element of token 0's row with a one-hot power of + two sized so the logit lands near +-8. Every product is exact, and each + output row is token 0's softmax weight for one code. A reference that + flushed K's 14 subnormal codes to zero moves their weights to 0.5 and is + rejected, so the gate sees a misread of a single code.""" + args, pages = _two_token_rows(kernel, seed=12100) + codes = _every_finite_e4m3_code() + keys, values = args["k_cache"].view(torch.uint8), args["v_cache"].view(torch.uint8) + keys[pages, :, 0, :] = codes + keys[pages, :, 1, :] = 0 + values[pages, :, 0, :] = torch.tensor(1.0, device=DEV).to(FP8).view(torch.uint8) + values[pages, :, 1, :] = 0 + q = args["q"] + rows, heads, dim = q.shape + assert rows * heads >= dim, "every code must be read by some (row, head)" + magnitude = codes.view(FP8).float().abs() + # |q . k| * softmax_scale = 128 / 16 = 8 up to the power-of-two rounding. + scale = torch.where(magnitude > 0, torch.exp2(torch.round(torch.log2(128.0 / magnitude))), 1.0) + column = (torch.arange(rows * heads, device=DEV) % dim).reshape(rows, heads, 1) + q.zero_() + q.scatter_(2, column, scale[column].to(q.dtype)) + out = op(**args) + assert launches[0][0] == kernel + _compare(out, op.reference(**args)) + + subnormal = ((codes & 0x78) == 0) & ((codes & 0x07) != 0) + assert int(subnormal.sum()) == 14 + flushed = args["k_cache"].clone() + flushed.view(torch.uint8)[pages, :, 0, :] = torch.where(subnormal, codes & 0x80, codes) + wrong = op.reference(**{**args, "k_cache": flushed}) + with pytest.raises(AssertionError): + _compare(out, wrong) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4.py new file mode 100644 index 000000000000..8c3c4a31e8a8 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4.py @@ -0,0 +1,416 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qsa_paged_sparse_gqa_nvfp4 entry. + +The paged state is real: one NVFP4 pool as the QSA cache manager serves it -- +e2m1 data `[pages, 2, kv, tpb, D/2]` and e4m3 block scales `[pages, 2, kv, tpb, +D/16]` -- sized from the pages its requests map, and the engine's slot table, +a row per request slot, more slots than requests, assigned by a seeded +permutation, `-1` past each sequence. Selection rows are built the way QSA's +top-k + expansion lay them out, packed at the front. + +Beside `_compare` against the reference, every cell runs QSA's bf16 sparse GQA +over a bf16 pool that holds the dequantized values and holds the two outputs to +the same band: the NVFP4 kernels are the bf16 ones with a dequantizing load. +""" + +import math +import re + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention import ( # noqa: E501 + qsa_paged_sparse_gqa_nvfp4 as _entry, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_sparse_gqa_nvfp4 import ( # noqa: E501 + qsa_paged_sparse_gqa_nvfp4 as op, +) +from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import triton_qsa_paged_sparse_gqa + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401 + +assert torch.cuda.is_available(), "qsa_paged_sparse_gqa_nvfp4 requires a CUDA device" + +DEV = "cuda" +ENGINE_MAX_SEQ_LEN = 8192 +#: The e4m3fn NaN bit pattern. +E4M3_NAN = 0x7F + +_KERNELS = { + "_nvfp4_sparse_gqa_prefill_kernel": "prefill", + "_nvfp4_sparse_gqa_splitk_kernel": "splitk", + "_nvfp4_merge_splitk_kernel": "merge", +} + + +class _LaunchRecorder: + """Records a Triton kernel's grid indexing, then delegates to the kernel.""" + + def __init__(self, inner, name, log): + self._inner, self._name, self._log = inner, name, log + + def __getitem__(self, grid): + self._log.append((self._name, tuple(grid))) + return self._inner[grid] + + +@pytest.fixture +def launches(monkeypatch): + """Which branch kernels the next call launches, in order, with grids.""" + log: list = [] + for name, tag in _KERNELS.items(): + monkeypatch.setattr(_entry, name, _LaunchRecorder(getattr(_entry, name), tag, log)) + return log + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """4 ulp of each (row, head)'s magnitude, 2 ulp relative RMS, in q's dtype. + + After the dequantize the kernels are the bf16 sparse GQA kernels, whose band + this is (they round the online-softmax probabilities to the value dtype + before P.V); the dequantize itself is exact at unit scale and rounded once, + as `reference` rounds it, otherwise. + """ + assert out.dtype == ref.dtype, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + if out.numel() == 0: + return + assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0) + + +def _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"] + segments = [ + (0 if i == 0 else (907 * i) % (ENGINE_MAX_SEQ_LEN - n + 1), n) + for i, n in enumerate(context_lens) + ] + 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: chosen complete groups, the trailing partial group, -1.""" + block_topk = token_topk // compress_ratio + final_topk = token_topk + compress_ratio - 1 + gen = torch.Generator(device=DEV).manual_seed(seed) + rows = positions.numel() + visible = positions.long() + 1 + available = visible // compress_ratio + taken = available.clamp(max=block_topk) + width = max(block_topk, int(available.max().item()) if rows else 0) + scores = torch.rand(rows, width, generator=gen, device=DEV) + group_ids = torch.arange(width, device=DEV)[None, :] + scores = torch.where(group_ids < available[:, None], scores, torch.inf) + chosen = scores.argsort(dim=1)[:, :block_topk] + column = torch.arange(final_topk, device=DEV)[None, :] + group_of_column = (column // compress_ratio).clamp(max=block_topk - 1).expand(rows, -1) + body = torch.gather(chosen, 1, group_of_column) * compress_ratio + column % compress_ratio + in_body = column < (taken * compress_ratio)[:, None] + tail_column = column - (taken * compress_ratio)[:, None] + tail = (available * compress_ratio)[:, None] + tail_column + in_tail = ~in_body & (tail_column >= 0) & (tail < visible[:, None]) + selected = torch.where(in_body, body, torch.where(in_tail, tail, torch.full_like(body, -1))) + return selected.to(torch.int32) + + +def _random_pool(pages, kv, tpb, d, gen): + """Pool bytes: any e2m1 nibble, and positive e4m3 scales from 2^-6 to 0.5.""" + data = torch.randint(0, 256, (pages, 2, kv, tpb, d // 2), generator=gen, device=DEV) + codes = torch.randint(0x08, 0x31, (pages, 2, kv, tpb, d // 16), generator=gen, device=DEV) + return data.to(torch.uint8), codes.to(torch.uint8) + + +def _build(spec, seed): + """Real paged state for one spec, keyed by the op's parameter names.""" + gen = torch.Generator(device=DEV).manual_seed(seed) + tpb, kv, d, q_heads = (spec[k] for k in ("tokens_per_block", "kv_heads", "head_dim", "q_heads")) + segments, 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)] + mapped = [math.ceil((s + n) / tpb) for s, n in segments] + assert not mapped or max(mapped) <= spec["pages"] + needed = sum(mapped) + 7 + slot_pages, layer = spec["slot_pages"], spec["layer"] + planes = slice(2 * layer, 2 * layer + 2) + # The requests sit on the pool's last `needed` pages, and every scale outside them is + # NaN, so a page address that strayed could not come back finite. + pool_pages = spec.get("pool_pages", needed) + assert pool_pages >= needed + first = pool_pages - needed + page_ids = first + torch.randperm(needed, generator=gen, device=DEV).to(torch.int32) + data = torch.zeros(pool_pages, slot_pages, kv, tpb, d // 2, dtype=torch.uint8, device=DEV) + scales = torch.full( + (pool_pages, slot_pages, kv, tpb, d // 16), E4M3_NAN, dtype=torch.uint8, device=DEV + ) + data[first:, planes], scales[first:, planes] = _random_pool(needed, kv, tpb, d, gen) + block_table = torch.full((slots, spec["pages"]), -1, dtype=torch.int32, device=DEV) + cursor = 0 + for slot, count in zip(slot_of.tolist(), mapped): + block_table[slot, :count] = page_ids[cursor : cursor + count] + cursor += count + + 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 + request_indices[2] = slots + 67 + block_table[:, 0] = -1 + + args = dict( + q=q, + k_cache=data[:, 2 * layer], + v_cache=data[:, 2 * layer + 1], + k_scale_cache=scales[:, 2 * layer], + v_scale_cache=scales[:, 2 * layer + 1], + kv_scale_quant_orig=torch.tensor( + [1.0, spec["global_k"], spec["global_v"]], dtype=torch.float32, device=DEV + ), + block_table=block_table, + selected_tokens=selected, + request_indices=request_indices, + tokens_per_block=tpb, + softmax_scale=d**-0.5, + ) + if spec["prefill_pair"]: + args.update(query_positions=positions, compress_ratio=spec["compress_ratio"]) + return args, generation + + +def _bf16_twin(args): + """QSA's bf16 sparse GQA over a bf16 pool holding the dequantized values.""" + pages = args["k_cache"].shape[0] + all_pages = torch.arange(pages, device=DEV)[:, None].expand(-1, args["tokens_per_block"]) + slots = torch.arange(args["tokens_per_block"], device=DEV)[None, :].expand(pages, -1) + planes = [] + for data, scales, index, grouped in ( + (args["k_cache"], args["k_scale_cache"], 1, False), + (args["v_cache"], args["v_scale_cache"], 2, True), + ): + values = _entry._dequantize_tokens( + data, scales, all_pages, slots, args["kv_scale_quant_orig"][index], grouped + ) # [pages, tpb, H, D] + planes.append(values.to(args["q"].dtype).transpose(1, 2).contiguous()) + twin = { + name: args[name] + for name in ("q", "block_table", "selected_tokens", "request_indices", "tokens_per_block") + } + twin.update(softmax_scale=args["softmax_scale"], k_cache=planes[0], v_cache=planes[1]) + if "query_positions" in args: + twin.update(query_positions=args["query_positions"], compress_ratio=args["compress_ratio"]) + return triton_qsa_paged_sparse_gqa(**twin) + + +def _bound(positions, spec): + cr = spec["compress_ratio"] + visible = positions.long() + 1 + return (visible // cr).clamp(max=spec["token_topk"] // cr) * cr + visible % cr + + +def _is_prefix_packed(table): + valid = table >= 0 + columns = torch.arange(table.shape[1], device=table.device)[None, :] + return (valid == (columns < valid.sum(1)[:, None])).all(1) + + +_OPERANDS = ( + "q", + "k_cache", + "v_cache", + "k_scale_cache", + "v_scale_cache", + "block_table", + "selected_tokens", + "request_indices", +) + + +def _capture_sweep(spec, seed, launches): + """The decode site under a CUDA graph, at each G for each rows-per-request: + capture, write a second state's operands in place, replay, and require the + eager call on that second state bitwise.""" + sweep = spec["capture"] + for per_request in sweep["rows_per_request"]: + for groups in sweep["groups"]: + _capture_once( + _decode_step(spec, groups, per_request), + seed + 100 * per_request + groups, + launches, + where=f"G={groups} x {per_request} rows", + ) + + +def _start(request: int, rows: int) -> int: + """A request's first position, across the budget.""" + return min(100 + (499 * request) % 8000, ENGINE_MAX_SEQ_LEN - rows) + + +def _decode_step(spec, groups, per_request): + """The decode site's step: `groups` requests x `per_request` rows.""" + segments = tuple((_start(g, per_request), per_request) for g in range(groups)) + return {**spec, "segments": segments} + + +def _capture_once(shape, seed, launches, where) -> None: + """One call of `shape` captured, a second state's operands written in place, + replayed: the replay must be bitwise the eager call on the second state.""" + where = f"{where}, [{shape['slots']}, {shape['pages']}] table" + args, _ = _build(shape, seed) + fresh, _ = _build(shape, seed + 50) + launches.clear() + op(**args) + assert [name for name, _ in launches] == list(shape["launches"]), where + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + try: + with torch.cuda.graph(graph): + captured = op(**args) + for name in _OPERANDS: + args[name].copy_(fresh[name]) + graph.replay() + torch.cuda.synchronize() + eager = op(**fresh) + assert torch.equal(captured, eager), f"replay differs from eager at {where}" + _compare(captured, op.reference(**fresh)) + finally: + del graph + torch.cuda.synchronize() + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying, launches) -> None: # noqa: F811 + certifying(op) + spec = cell.spec + if "capture" in spec: + _capture_sweep(spec, seed=sum(map(ord, cell.why)), launches=launches) + return + args, _ = _build(spec, seed=sum(map(ord, cell.why))) + rows = args["q"].shape[0] + if spec["prefill_pair"] and not spec.get("holes"): + selected = args["selected_tokens"] + assert bool(_is_prefix_packed(selected).all()) + valid = (selected >= 0).sum(1) + assert torch.equal(valid, _bound(args["query_positions"], spec)) + + out = op(**args) + assert [name for name, _ in launches] == list(spec["launches"]) + assert launches[0][1][:2] == (rows, spec["kv_heads"]) + assert out.shape == (rows, spec["q_heads"], spec["head_dim"]) + assert out.dtype == args["q"].dtype and out.is_contiguous() + assert not bool(out.isnan().any()) + _compare(out, op.reference(**args)) + # QSA's bf16 kernels over the same values: inside the same band, not bitwise + # (the dequantizing load changes the compiled kernel). + _compare(out, _bf16_twin(args)) + + +CKPT = op.CELLS[0].spec + + +def _with(**overrides): + return {**CKPT, **overrides} + + +def test_gate_discriminates_wrong_variants() -> None: + """References wrong one way each -- V's scales read token-major like K's, + the nibbles swapped, the second-level scale ignored, K and V swapped -- are + all rejected by `compare`.""" + args, _ = _build(_with(segments=((2032, 272),), prefill_pair=True, global_v=2.5), seed=9000) + out = op(**args) + _compare(out, op.reference(**args)) + v_scales = args["v_scale_cache"] + pages, kv, tpb, groups = v_scales.shape + token_major = ( + v_scales.reshape(pages, kv, tpb // 4, 4, groups) + .permute(0, 1, 2, 4, 3) + .reshape(v_scales.shape) + ) + nibbles = (args["v_cache"] >> 4) | (args["v_cache"] << 4) + variants = { + "v_scales_token_major": dict(v_scale_cache=token_major), + "nibbles_swapped": dict(v_cache=nibbles), + "no_second_level_scale": dict( + kv_scale_quant_orig=torch.ones(3, dtype=torch.float32, device=DEV) + ), + "k_and_v_swapped": dict( + k_cache=args["v_cache"], + v_cache=args["k_cache"], + k_scale_cache=args["v_scale_cache"], + v_scale_cache=args["k_scale_cache"], + ), + } + for name, override in variants.items(): + wrong = op.reference(**{**args, **override}) + with pytest.raises(AssertionError): + _compare(out, wrong) + assert (out.float() - wrong.float()).abs().max().item() > 0.05, name + + +def test_calls_are_bitwise_deterministic() -> None: + args, _ = _build(_with(segments=tuple((4000 + 37 * i, 2) for i in range(4))), seed=8800) + assert torch.equal(op(**args), op(**args)) + + +def test_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)) + + +_REJECTIONS = { + "float32_query": (lambda a: {**a, "q": a["q"].float()}, "q must be BF16 or FP16"), + "bf16_pool": ( + lambda a: {**a, "k_cache": a["k_cache"].view(torch.int8)}, + "k_cache must be uint8", + ), + "scales_of_another_width": ( + lambda a: {**a, "v_scale_cache": a["v_scale_cache"][..., :8]}, + "v_scale_cache must be uint8", + ), + "single_scale": ( + lambda a: {**a, "kv_scale_quant_orig": a["kv_scale_quant_orig"][:1]}, + "kv_scale_quant_orig must be a contiguous fp32 [3]", + ), + "int64_block_table": ( + lambda a: {**a, "block_table": a["block_table"].long()}, + "block_table must be a row-contiguous 2-D int32 tensor", + ), + "positions_without_ratio": ( + lambda a: {**a, "query_positions": torch.zeros(4, dtype=torch.int64, device=DEV)}, + "pass both query_positions and compress_ratio, or neither", + ), + "host_query": (lambda a: {**a, "q": a["q"].cpu()}, "every operand must be on q's CUDA device"), +} + + +@pytest.mark.parametrize("case", list(_REJECTIONS), ids=list(_REJECTIONS)) +def test_rejects(case) -> None: + """The op's own loud refusals -- the reason the entry carries no `is_valid`.""" + 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)): + with validating(op): + op(**mutate(args)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_prefill_compress.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_prefill_compress.py index 05b1b5a03d10..0999b98a150c 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_prefill_compress.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_prefill_compress.py @@ -28,24 +28,37 @@ assert torch.cuda.is_available(), "qsa_prefill_compress requires a CUDA device" -def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: - """Bitwise over the whole cache. - - Measured on sm_103: the kernel matches this fp32 reference on every - element of every cell, so there is no tolerance to derive. It has to - be bitwise -- a pool accumulated in bf16 instead of fp32 lands only - 0.03125 away, inside any bf16 band. Comparing every slot, not just the - anchors, is the untouched-slot lock. +def _compare(out: torch.Tensor, ref: torch.Tensor, flips: int = 0) -> None: + """Bitwise over the whole cache, except for `flips` elements one bf16 step off. + + Measured on sm_103: the kernel matches this fp32 reference on every element of every + cell but one, so there is no tolerance to derive. It has to be bitwise -- a pool + accumulated in bf16 instead of fp32 lands only 0.03125 away, inside any bf16 band. + Comparing every slot, not just the anchors, is the untouched-slot lock. + + The exception is the cell of 32 MTP7 verify blocks (64 anchors, `rounding_flips`): the + kernel takes `tl.rsqrt`, the hardware approximation, where the reference takes torch's, + and the two differ by an fp32 ulp, which moves the bf16 rounding of an element whose fp32 + value sits on a rounding boundary. One element of its 8192 anchor elements is one bf16 + step off. """ assert out.dtype == ref.dtype, (out.dtype, ref.dtype) assert out.shape == ref.shape, (out.shape, ref.shape) - torch.testing.assert_close(out.view(torch.int16), ref.view(torch.int16), rtol=0, atol=0) + if flips == 0: + torch.testing.assert_close(out.view(torch.int16), ref.view(torch.int16), rtol=0, atol=0) + return + distance = (out.view(torch.int16).int() - ref.view(torch.int16).int()).abs() + assert int(distance.max()) <= 1 and int((distance != 0).sum()) <= flips, ( + f"{int((distance != 0).sum())} elements differ, by up to {int(distance.max())} bf16 " + f"steps; at most {flips} may differ by one" + ) DEV = "cuda" ROTARY_DIM = 64 THETA = 1.0e7 MAX_POSITIONS = 16384 +SENTINEL = 0x3C3C def _cos_sin() -> torch.Tensor: @@ -103,6 +116,24 @@ def _build(spec, seed: int): pool[pages:] = index_cache index_cache = pool[pages:] extras["neighbour"] = (pool[:pages], pool[:pages].clone()) + elif spec["layout"] == "shared_page": + # An NVFP4 pool's page: `roles` index pages, then the position page. The other + # index roles hold a sentinel and must survive, as must the position page. + roles = spec["roles"] + index_bytes = tpb * head_dim * 2 + page = torch.empty(pages, roles * index_bytes + tpb * 3 * 4, dtype=torch.uint8, device=DEV) + others = page[:, : (roles - 1) * index_bytes].view(torch.int16) + others.fill_(SENTINEL) + shared_index = page[:, (roles - 1) * index_bytes : roles * index_bytes] + shared_index = shared_index.view(torch.bfloat16).view(pages, tpb, 1, head_dim) + shared_position = page[:, roles * index_bytes :].view(torch.int32).view(pages, tpb, 3) + shared_index.copy_(index_cache) + shared_position.copy_(position_cache) + assert shared_index.stride(0) == roles * tpb * head_dim + tpb * 3 * 2 + assert shared_position.stride() == (roles * tpb * head_dim // 2 + tpb * 3, 3, 1) + index_cache, position_cache = shared_index, shared_position + extras["neighbour"] = (others, others.clone()) + extras["position_page"] = (position_cache, position_cache.clone()) elif spec["layout"] == "head_dim_strided": wide = torch.randn(pages, tpb, 1, 2 * head_dim, device=DEV).to(torch.bfloat16) wide[..., ::2] = index_cache @@ -140,12 +171,15 @@ def test_certified_cells(cell, certifying): # noqa: F811 ref = op.reference(**args) out = op(**args) assert out is args["index_cache"] - _compare(out, ref) + _compare(out, ref, cell.spec.get("rounding_flips", 0)) changed = (out.view(torch.int16) != before.view(torch.int16)).any(-1).sum() assert int(changed) == _anchors(args), "exactly the anchor slots are written" if "neighbour" in extras: view, saved = extras["neighbour"] assert torch.equal(view, saved), "wrote outside the strided view" + if "position_page" in extras: + view, saved = extras["position_page"] + assert torch.equal(view, saved), "wrote into the position page" def _scene(seed: int, **overrides): @@ -172,9 +206,15 @@ def test_a_second_call_over_the_same_rows_is_not_idempotent() -> None: assert int(changed) == _anchors(args) -def test_cuda_graph_capture_and_replay() -> None: +#: The cell of 32 MTP7 verify blocks (256 rows, two anchors each): the segments, pool and table. +_MTP7_32 = next(c.spec for c in op.CELLS if c.why.startswith("decode MTP7: 32 requests")) +_MTP7_BATCH = {key: _MTP7_32[key] for key in ("segments", "pages", "slots", "rounding_flips")} + + +@pytest.mark.parametrize("batch", ({}, _MTP7_BATCH), ids=("prefill", "mtp7-32x8")) +def test_cuda_graph_capture_and_replay(batch) -> None: """The launch grid is fixed by the row count, so it captures and replays.""" - args = _scene(seed=16) + args = _scene(seed=16, **batch) cache = args["index_cache"] before = cache.clone() ref = op.reference(**args) @@ -194,7 +234,7 @@ def test_cuda_graph_capture_and_replay() -> None: cache.copy_(before) graph.replay() torch.cuda.synchronize() - _compare(cache, ref) + _compare(cache, ref, batch.get("rounding_flips", 0)) def test_stream_ordering_needs_no_explicit_synchronization() -> None: diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_select.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_select.py index 11421ec14395..b687789a1663 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_select.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_select.py @@ -140,10 +140,12 @@ def test_nan_inside_the_window_is_selected() -> None: 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.""" +@pytest.mark.parametrize("rows", (8, 256)) +def test_cuda_graph_capture_and_replay(rows) -> None: + """Captured, then replayed against new scores and a new causal bound: 8 rows, and the + 256 rows a verify step of 32 requests carries at max_draft_len 7.""" gen = torch.Generator(device=DEV).manual_seed(8) - rows, cols = 8, 2048 + cols = 2048 logits = _logits("relu", rows, cols, gen) row_starts = torch.zeros(rows, dtype=torch.int32, device=DEV) row_ends = torch.full((rows,), cols, dtype=torch.int32, device=DEV) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py new file mode 100644 index 000000000000..df6d50d69c22 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py @@ -0,0 +1,711 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the CuteDSL NVFP4 FC1 entry of the Qwen3.8-Flash-Next target. + +The entry owns the contract. This file builds the checkpoint-shaped operands -- +512 experts in the CuteDSL weight layout, the activation as the target quantizes +it, the sort as moe_sort returns it -- and holds the claims no cell carries: +per-pair invariance across tiles, tactics, row order and prefixes; CUDA-graph +replay at the captured row counts; the regions the op leaves unwritten; that +the gate rejects every wrong layout; the guards and the loud rejections. +""" + +import functools +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion import ( # noqa: E501 + E2M1, + E2M1_MAX, + GATE_UP_GROUP, + Fc1Expected, + e2m1_codes, + e4m3_value, + unswizzle, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion import ( # noqa: E501 + cute_dsl_nvfp4_gather_grouped_gemm_act_fusion as op, +) +from tensorrt_llm._torch.autotuner import AutoTuner +from tensorrt_llm._torch.utils import ActivationType + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "the CuteDSL FC1 requires a CUDA device" + +DEV = "cuda" +E, K, H, INTER, VEC = 512, 10, 2560, 640, 16 +SWIGLU = int(ActivationType.Swiglu) +#: The activation's input scale: |x| = 5 maps to the e4m3 maximum, as the +#: checkpoint's per-layer `input_scale` does for its own calibration maximum. +ISC1 = 5.0 / (448 * 6) +#: The SwiGLU output maximum `global_sf` is calibrated to (see `_ckpt` in the entry). +ACT_AMAX = 48.0 +GSF = { + "calibrated": 448 * 6 / ACT_AMAX, + "saturating": 25 * 448 * 6 / ACT_AMAX, + "subnormal": 2.0**-12 * 448 * 6 / ACT_AMAX, +} +SENTINEL = 0xA5 + +#: Tie windows of `_compare`, derived in its docstring. +SF_TIE = 2.0**-20 +CODE_TIE = 2.0**-17 +#: At most this fraction of block scales / codes may take the other side of a tie +#: (4 of them when that is more). +FLIP_BUDGET = 1e-5 + +#: Row counts of a decode capture (requests x MTP verify width 1..4, up to 32 +#: requests), and the eager prefix lengths: both sides of the power-of-two edges up +#: to the 8192-row ceiling. +_CAPTURED_ROWS = (1, 2, 3, 4, 5, 8, 32, 96, 128) +_PREFIXES = (1, 2, 3, 4, 5, 8, 31, 32, 33, 127, 128, 129, 512, 1023, 1024, 1025, 4095, 4096, 8191) + + +# ── the gate ────────────────────────────────────────────────────────────── + + +def _e2m1_rank(packed: torch.Tensor) -> torch.Tensor: + """Packed codes -> signed magnitude rank in -7..7 per element; -0 and +0 are both 0.""" + lo, hi = (packed & 0xF).long(), (packed >> 4).long() + codes = torch.stack([lo, hi], dim=-1).reshape(packed.shape[0], -1) + return torch.where(codes >= 8, -(codes & 7), codes & 7) + + +def _e2m1_rank_of(v: torch.Tensor) -> torch.Tensor: + codes = e2m1_codes(v).long() + return torch.where(codes >= 8, -(codes & 7), codes & 7) + + +def _rank_value(rank: torch.Tensor) -> torch.Tensor: + lut = torch.tensor(E2M1, dtype=torch.float64, device=rank.device) + return lut[rank.abs()] * torch.sign(rank).double() + + +def _deviation(out: tuple[torch.Tensor, torch.Tensor], ref: Fc1Expected) -> dict: + """Every disagreement between the op and the reference on the rows it writes. + + Block scales: equal, or the two neighbours of an e4m3 midpoint that the exact + pre-rounding scale lies within `SF_TIE` (relative) of -- a `sf_flip`; anything + else is `sf_outside`. Codes are re-derived under the op's OWN block scale (so + a flipped scale does not cascade into its 16 codes): equal, or the two + neighbours of an e2m1 midpoint that the exact `scaled / sf` lies within + `CODE_TIE` of -- a `code_flip`; anything else is `code_outside`. Blocks whose + scale is zero dequantize to zero whatever their codes hold and are not + code-checked. `*_far` counts the disagreements that are not between + neighbours at all. + """ + c, c_sf = out + m, half = c.shape + inter = half * 2 + nb = c_sf.numel() // m + vec = inter // nb + rows = ref.rows + s_op = unswizzle(c_sf.view(torch.uint8), m, nb)[rows].long() + s_ref = unswizzle(ref.c_sf.view(torch.uint8), m, nb)[rows].long() + pre = ref.scaled.view(-1, nb, vec).abs().amax(-1) / E2M1_MAX + + v_op, v_ref = e4m3_value(s_op), e4m3_value(s_ref) + sf_diff = s_op != s_ref + finite = (s_op < 0x7F) & (s_ref < 0x7F) + sf_adj = sf_diff & finite & ((s_op - s_ref).abs() == 1) + sf_dist = ((pre - (v_op + v_ref) / 2).abs() / pre.clamp_min(1e-300)).masked_fill(~sf_adj, 0.0) + sf_flip = sf_adj & (sf_dist <= SF_TIE) + sf_outside = sf_diff & ~sf_flip + + rank_op = _e2m1_rank(c.view(torch.uint8)[rows]) + scale = v_op.repeat_interleave(vec, dim=1) + live = scale > 0 + q = torch.where(live, ref.scaled / scale, torch.zeros_like(scale)) + rank_exp = _e2m1_rank_of(q) + code_diff = live & (rank_op != rank_exp) + code_adj = code_diff & ((rank_op - rank_exp).abs() == 1) + mid = (_rank_value(rank_op) + _rank_value(rank_exp)) / 2 + code_dist = (q - mid).abs().masked_fill(~code_adj, 0.0) + code_flip = code_adj & (code_dist <= CODE_TIE) + code_outside = code_diff & ~code_flip + + return dict( + rows=int(rows.sum()), + blocks=s_op.numel(), + sf_flips=int(sf_flip.sum()), + sf_outside=int(sf_outside.sum()), + sf_far=int((sf_diff & ~sf_adj).sum()), + elements=int(live.sum()), + code_flips=int(code_flip.sum()), + code_outside=int(code_outside.sum()), + code_far=int((code_diff & ~code_adj).sum()), + ) + + +def _compare(out: tuple[torch.Tensor, torch.Tensor], ref: Fc1Expected) -> None: + """Bitwise on every written row, except at rounding ties of bounded width. + + The op accumulates in fp32 on the tensor cores, evaluates SwiGLU with + `ex2.approx` / `rcp.approx`, and divides by the block scale through + `rcp.approx`; the reference is exact. The op's pre-rounding value is + therefore off by a handful of fp32 roundings (2^-24 relative each), and + only where the exact value lies that close to a rounding midpoint may the + op land on its other side. Everywhere else the two agree bit for bit. + + - Block scale: a flip only within `SF_TIE` = 2^-20 (16 fp32 roundings) + relative of an e4m3 midpoint. + - Code: a flip only within `CODE_TIE` = 2^-17 of an e2m1 midpoint, in units + of the block scale -- 16 fp32 roundings of a block maximum of at most 6, + rounded up to a power of two. + - Flips: at most `FLIP_BUDGET` = 1e-5 of the scales and of the codes (4 when + that is fewer), so the op may not take the far side of every near-tie. + + Every wrongly prepared operand puts most of the block scales outside the gate. + """ + c, c_sf = out + assert c.dtype == torch.float4_e2m1fn_x2, c.dtype + assert c_sf.dtype == torch.uint8, c_sf.dtype + assert c.shape == ref.c.shape, (tuple(c.shape), tuple(ref.c.shape)) + assert c_sf.shape == ref.c_sf.shape, (tuple(c_sf.shape), tuple(ref.c_sf.shape)) + d = _deviation(out, ref) + assert d["rows"] > 0, d + assert d["sf_outside"] == 0, f"block scales off by more than a tie: {d}" + assert d["code_outside"] == 0, f"codes off by more than a tie: {d}" + assert d["sf_flips"] <= max(4, FLIP_BUDGET * d["blocks"]), f"too many scale ties flipped: {d}" + assert d["code_flips"] <= max(4, FLIP_BUDGET * d["elements"]), ( + f"too many code ties flipped: {d}" + ) + + +# ── operands ────────────────────────────────────────────────────────────── + + +@functools.lru_cache(maxsize=None) +def _experts(num_experts: int = E): + """Checkpoint-format routed experts: per expert `[up; gate]` NVFP4 rows with + LINEAR e4m3 block scales and one `weight_scale_2` shared by both halves, as the + checkpoint stores them. Expert magnitudes spread over 2^-1..2^1 so a wrong + per-expert alpha shows.""" + g = torch.Generator(device=DEV).manual_seed(20261013) + spread = torch.exp2(torch.rand(num_experts, generator=g, device=DEV) * 2 - 1) + w = torch.empty(num_experts, 2 * INTER, H // 2, dtype=torch.uint8, device=DEV) + s = torch.empty(num_experts, 2 * INTER, H // VEC, dtype=torch.uint8, device=DEV) + ws2 = torch.empty(num_experts, dtype=torch.float32, device=DEV) + for e in range(num_experts): + dense = torch.randn(2 * INTER, H, generator=g, device=DEV) * (float(spread[e]) / H**0.5) + dense = dense.to(torch.bfloat16) + gs = (448 * 6 / dense.abs().amax().float()).reshape(1) + q, sf = torch.ops.trtllm.fp4_quantize(dense, gs, VEC, False, False) + w[e] = q.view(2 * INTER, H // 2) + s[e] = sf.view(2 * INTER, H // VEC) + ws2[e] = 1.0 / gs + return w, s, ws2 + + +def _row_perm(n: int, layout: str) -> torch.Tensor: + """Destination row -> source row of the `[up; gate]` stack.""" + idx = torch.arange(n, device=DEV) + if layout == "group64": + return idx.view(2, n // (2 * GATE_UP_GROUP), GATE_UP_GROUP).transpose(0, 1).reshape(-1) + if layout == "rowwise": # trtllm-gen's interleave, without its 32-row shuffle + return idx.view(2, n // 2).transpose(0, 1).reshape(-1) + if layout == "gate_first": + half = idx.view(2, n // 2).flip(0).reshape(-1) + return half[_row_perm(n, "group64")] + raise ValueError(layout) + + +@functools.lru_cache(maxsize=None) +def _weights(layout: str = "group64", num_experts: int = E): + """`(weight, weight_scale, alpha)` in the layout CuteDSL reads -- what the + target's load-time transform produces -- or one of the wrong layouts the + discrimination test feeds it.""" + w, s, ws2 = _experts(num_experts) + perm = _row_perm( + 2 * INTER, "group64" if layout in ("linear_scales", "scales_unpermuted") else layout + ) + w_k = w[:, perm].contiguous() + s_rows = s if layout == "scales_unpermuted" else s[:, perm] + if layout == "linear_scales": + s_k = s_rows.contiguous() + else: + s_k = torch.ops.trtllm.block_scale_interleave(s_rows.contiguous()).view(s.shape) + alpha = (ISC1 * ws2).contiguous() + return w_k.view(torch.float4_e2m1fn_x2), s_k, alpha + + +def _activation(tokens: int, seed: int, zero_tokens: int = 0, swizzled: bool = False): + """The target's FC1 input: bf16 hidden rows through `fp4_quantize(x, 1/isc1, + 16, False, False)`, scales viewed `[T, H/16]`.""" + g = torch.Generator(device=DEV).manual_seed(seed) + x = torch.randn(tokens, H, generator=g, device=DEV).to(torch.bfloat16) + x[:zero_tokens] = 0 + gs = torch.tensor([1.0 / ISC1], dtype=torch.float32, device=DEV) + q, sf = torch.ops.trtllm.fp4_quantize(x, gs, VEC, False, swizzled) + return q.view(torch.float4_e2m1fn_x2), sf.view(tokens, H // VEC) + + +def _ids(dist: str, tokens: int, tile: int, seed: int, num_experts: int = E) -> torch.Tensor: + """Expert ids `[tokens, K]` int32, every row naming K distinct experts.""" + g = torch.Generator(device=DEV).manual_seed(seed) + if dist == "router": + scores = torch.rand(tokens, num_experts, device=DEV, generator=g) + return scores.topk(K, dim=-1).indices.to(torch.int32) + if dist == "hot": + hot = torch.randperm(num_experts, device=DEV, generator=g)[:K] + slots = torch.rand(tokens, K, device=DEV, generator=g).argsort(dim=-1) + return hot[slots].to(torch.int32) + if dist == "tile_edges": + edge = (7, 300, 511) + scores = torch.rand(tokens, num_experts, device=DEV, generator=g) + scores[:, list(edge)] = -1.0 + ids = scores.topk(K, dim=-1).indices.to(torch.int32) + for slot, (expert, rows) in enumerate(zip(edge, (tile - 1, tile, tile + 1))): + ids[:rows, slot] = expert + return ids + if dist == "round_robin": + base = torch.arange(tokens, device=DEV).unsqueeze(1) * K + return ((base + torch.arange(K, device=DEV)) % num_experts).to(torch.int32) + raise ValueError(dist) + + +def _sort(ids: torch.Tensor, tile: int, num_experts: int = E): + """moe_sort as the CuteDSL MoE calls it; the routing weights do not reach FC1.""" + scales = torch.full(ids.shape, 0.1, dtype=torch.float32, device=DEV) + return torch.ops.trtllm.moe_sort( + token_selected_experts=ids, + token_final_scales=scales, + num_experts=num_experts, + top_k=K, + local_expert_offset=0, + local_num_experts=num_experts, + tile_tokens_dim=tile, + ) + + +def _args(x, x_sf, sort, tile, gsf, layout="group64"): + """The keyword arguments the target passes, by the op's own names.""" + weight, weight_scale, alpha = _weights(layout) + t2e, t2mn, _e2p, p2e, _total, n_tiles = sort + return dict( + input=x, + weight=weight, + input_scale=x_sf, + weight_scale=weight_scale, + alpha=alpha, + tile_idx_to_group_idx=t2e, + tile_idx_to_mn_limit=t2mn, + permuted_idx_to_expanded_idx=p2e, + num_non_exiting_tiles=n_tiles, + global_sf=torch.tensor([gsf], dtype=torch.float32, device=DEV), + num_experts=E, + top_k=K, + num_local_experts=E, + local_expert_offset=0, + tile_size=tile, + scaling_vector_size=VEC, + activation_type=SWIGLU, + ) + + +def _problem(tokens, tile, seed, dist="router", gsf="calibrated", zero_tokens=0): + x, x_sf = _activation(tokens, seed, zero_tokens) + sort = _sort(_ids(dist, tokens, tile, seed + 1), tile) + return _args(x, x_sf, sort, tile, GSF[gsf]), sort + + +def _call(args, mma_n): + """One op call with the autotuner tactic whose MMA tile is `(tile_size, mma_n)`, + selected through the tuner's own capture/replay so the call still goes through + the op.""" + tuner = AutoTuner.get() + with tuner.capture() as cap: + op(**args) + for config in cap: + ((_runner, tactic),) = config + if isinstance(tactic, tuple) and tactic[0] == (args["tile_size"], mma_n): + with tuner.replay(config): + return op(**args) + raise AssertionError(f"no tactic with MMA tile ({args['tile_size']}, {mma_n})") + + +def _pairs(out, sort, tokens): + """Each (token, slot) pair's output row and its block-scale row, in expanded + order -- what FC2 and the combine see of it, whatever row the sort gave it.""" + c, c_sf = out + e2p = sort[2].reshape(-1).long()[: tokens * K] + sf = unswizzle(c_sf, c.shape[0], c_sf.numel() // c.shape[0]) + return c.view(torch.uint8)[e2p], sf[e2p] + + +def _clone(v): + """A copy that does not depend on copy kernels for the fp4 dtype.""" + if not isinstance(v, torch.Tensor): + return v + return v.view(torch.uint8).clone().view(v.dtype) if v.element_size() == 1 else v.clone() + + +def _assert_pairs_equal(a, b, what): + assert torch.equal(a[0], b[0]), f"{what}: codes differ" + assert torch.equal(a[1], b[1]), f"{what}: block scales differ" + + +def _cell_id(cell): + s = cell.spec + return ( + f"t{s['tile_size']}-n{s['mma_n']}-T{s['tokens']}-{s['dist']}-{s['gsf']}-z{s['zero_tokens']}" + ) + + +# ── the certified cells ─────────────────────────────────────────────────── + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[_cell_id(c) for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + s = cell.spec + assert (s["num_experts"], s["top_k"], s["hidden"], s["intermediate"]) == (E, K, H, INTER) + seed = zlib.crc32(cell.why.encode()) + args, sort = _problem(s["tokens"], s["tile_size"], seed, s["dist"], s["gsf"], s["zero_tokens"]) + before = {k: _clone(v) for k, v in args.items() if isinstance(v, torch.Tensor)} + out = _call(args, s["mma_n"]) + _compare(out, op.reference(**args)) + for k, v in before.items(): + assert torch.equal(args[k].view(torch.uint8), v.view(torch.uint8)), f"{k} was written" + if s["zero_tokens"]: + zero = sort[2][: s["zero_tokens"]].reshape(-1).long() + c, c_sf = out + sf = unswizzle(c_sf, c.shape[0], INTER // VEC) + assert bool((sf[zero] == 0).all()), "an all-zero token's block scales are not zero" + assert bool((c.view(torch.uint8)[zero] & 0x77 == 0).all()), "a zero token's codes are not 0" + + +# ── invariance: a pair's output depends on the pair alone ───────────────── + + +@pytest.mark.parametrize("tokens", (1, 128, 8192)) +def test_tiles_and_tactics_agree_per_pair(tokens, certifying): # noqa: F811 + """Tile 128 / 256, MMA N 128 / 256 and the tuner's fallback all produce every + pair's codes and block scales bit for bit the same, so a cell certifies the + tactic the autotuner picks in serving as well as the one it ran.""" + certifying(op) + x, x_sf = _activation(tokens, seed=tokens) + ids = _ids("router", tokens, 128, seed=tokens + 1) + seen = {} + for tile in (128, 256): + sort = _sort(ids, tile) + args = _args(x, x_sf, sort, tile, GSF["calibrated"]) + for mma_n in (128, 256): + seen[(tile, mma_n)] = _pairs(_call(args, mma_n), sort, tokens) + seen[(tile, "fallback")] = _pairs(op(**args), sort, tokens) + if tile == 128: + _compare(op(**args), op.reference(**args)) + first = seen[(128, 128)] + for key, got in seen.items(): + _assert_pairs_equal(got, first, f"T={tokens} {key} vs (128, 128)") + + +@pytest.mark.parametrize("tile", (128, 256)) +def test_row_order_and_repeat_do_not_change_a_pair(tile, certifying): # noqa: F811 + """moe_sort orders an expert's rows by atomics (from T = 5). Shuffling every + expert's run of `permuted_idx_to_expanded_idx`, and calling twice, leaves each + pair's output bitwise unchanged: the FC1 result is reproducible run to run.""" + certifying(op) + tokens = 4096 + args, sort = _problem(tokens, tile, seed=77) + base = _pairs(op(**args), sort, tokens) + _assert_pairs_equal(_pairs(op(**args), sort, tokens), base, "repeat call") + + t2mn, e2p, p2e, n_tiles = sort[1], sort[2], sort[3], int(sort[5]) + g = torch.Generator(device=DEV).manual_seed(5) + p2e_s = p2e.clone() + starts = torch.arange(n_tiles, device=DEV) * tile + run_of = torch.full((p2e.numel(),), -1, dtype=torch.long, device=DEV) + for i, (start, limit) in enumerate(zip(starts.tolist(), t2mn[:n_tiles].tolist())): + run_of[start:limit] = int(sort[0][i]) + for e in torch.unique(run_of[run_of >= 0]).tolist(): + rows = torch.nonzero(run_of == e).squeeze(1) + p2e_s[rows] = p2e[rows[torch.randperm(rows.numel(), generator=g, device=DEV)]] + e2p_s = e2p.clone().reshape(-1) + real = run_of >= 0 + e2p_s[p2e_s[real].long()] = torch.nonzero(real).squeeze(1).to(torch.int32) + shuffled = (sort[0], sort[1], e2p_s.view_as(e2p), p2e_s, sort[4], sort[5]) + args_s = dict(args, permuted_idx_to_expanded_idx=p2e_s) + out_s = op(**args_s) + _compare(out_s, op.reference(**args_s)) + _assert_pairs_equal(_pairs(out_s, shuffled, tokens), base, "shuffled rows") + + +@pytest.mark.parametrize("tile", (128, 256)) +def test_eager_rows_prefix_bitwise(tile, certifying): # noqa: F811 + """The 8192-row call is gated against the reference, and every prefix call's + pairs are bitwise the 8192-row call's.""" + certifying(op) + tokens = 8192 + x, x_sf = _activation(tokens, seed=8) + ids = _ids("router", tokens, tile, seed=9) + sort = _sort(ids, tile) + args = _args(x, x_sf, sort, tile, GSF["calibrated"]) + full = op(**args) + _compare(full, op.reference(**args)) + full_pairs = _pairs(full, sort, tokens) + for t in _PREFIXES: + sub_sort = _sort(ids[:t].contiguous(), tile) + sub = op(**_args(x[:t], x_sf[:t], sub_sort, tile, GSF["calibrated"])) + got = _pairs(sub, sub_sort, t) + n = t * K + _assert_pairs_equal(got, (full_pairs[0][:n], full_pairs[1][:n]), f"prefix T={t}") + + +# ── CUDA graphs: captured row counts, on new data ──────────────────────── + + +def _capture(fn): + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for _ in range(2): + fn() + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + static = fn() + return graph, static + + +@pytest.mark.parametrize("tile", (128, 256)) +def test_captured_rows_replay_on_new_data(tile, certifying): # noqa: F811 + """moe_sort and this op are captured together, as the target's graph holds them; + each graph is replayed on two fresh activations and id sets, and its written + rows are gated against the reference and bitwise against an eager call on the + graph's own sort output.""" + certifying(op) + for tokens in _CAPTURED_ROWS: + x, x_sf = _activation(tokens, seed=tokens) + ids = _ids("router", tokens, tile, seed=500 + tokens) + scales = torch.full(ids.shape, 0.1, dtype=torch.float32, device=DEV) + base = _args(x, x_sf, _sort(ids, tile), tile, GSF["calibrated"]) + state = {} + + def step(ids=ids, scales=scales, base=base, state=state, tile=tile): + sort = torch.ops.trtllm.moe_sort( + token_selected_experts=ids, + token_final_scales=scales, + num_experts=E, + top_k=K, + local_expert_offset=0, + local_num_experts=E, + tile_tokens_dim=tile, + ) + state["sort"] = sort + return op( + **dict( + base, + tile_idx_to_group_idx=sort[0], + tile_idx_to_mn_limit=sort[1], + permuted_idx_to_expanded_idx=sort[3], + num_non_exiting_tiles=sort[5], + ) + ) + + graph, static_out = _capture(step) + static_sort = state["sort"] + for trial in (1, 2): + nx, nx_sf = _activation(tokens, seed=1000 * trial + tokens) + x.view(torch.uint8).copy_(nx.view(torch.uint8)) + x_sf.copy_(nx_sf) + ids.copy_(_ids("router", tokens, tile, seed=7000 * trial + tokens)) + graph.replay() + args = dict( + base, + tile_idx_to_group_idx=static_sort[0], + tile_idx_to_mn_limit=static_sort[1], + permuted_idx_to_expanded_idx=static_sort[3], + num_non_exiting_tiles=static_sort[5], + ) + ref = op.reference(**args) + _compare(static_out, ref) + eager = op(**{k: _clone(v) for k, v in args.items()}) + _compare(eager, ref) + rows = ref.rows + assert torch.equal( + static_out[0].view(torch.uint8)[rows], eager[0].view(torch.uint8)[rows] + ), (tokens, trial) + m = rows.numel() + sf_g = unswizzle(static_out[1], m, INTER // VEC)[rows] + sf_e = unswizzle(eager[1], m, INTER // VEC)[rows] + assert torch.equal(sf_g, sf_e), (tokens, trial) + del graph, static_out + torch.cuda.synchronize() + + +@pytest.mark.parametrize("tile", (128, 256)) +@pytest.mark.parametrize("tokens", (3, 300)) +def test_unwritten_regions(tokens, tile): + """A sentinel planted in a captured graph's outputs: which bytes replay writes. + The rows the sort filled hold the result; the padding rows of a written tile + are written with zero block scales and zero codes; the tiles past + `num_non_exiting_tiles` are not written -- under a graph they keep whatever + the buffer held.""" + args, sort = _problem(tokens, tile, seed=11) + graph, static_out = _capture(lambda: op(**args)) + c, c_sf = static_out + m = c.shape[0] + for trial in (1, 2): + nx, nx_sf = _activation(tokens, seed=60 + trial) + args["input"].view(torch.uint8).copy_(nx.view(torch.uint8)) + args["input_scale"].copy_(nx_sf) + c.view(torch.uint8).fill_(SENTINEL) + c_sf.fill_(SENTINEL) + graph.replay() + ref = op.reference(**args) + _compare(static_out, ref) + n_tiles = int(sort[5]) + span = n_tiles * tile + codes = c.view(torch.uint8) + sf = unswizzle(c_sf, m, INTER // VEC) + assert bool((codes[span:] == SENTINEL).all()) and bool((sf[span:] == SENTINEL).all()), ( + "a tile past num_non_exiting_tiles was written" + ) + pad = ~ref.rows[:span] + assert int(pad.sum()) > 0 + assert bool((sf[:span][pad] == 0).all()), "a padding row's block scales are not 0" + assert bool((codes[:span][pad] == 0).all()), "a padding row's codes are not 0" + del graph + + +# ── the gate sees a wrong layout ────────────────────────────────────────── + + +@pytest.mark.parametrize( + "layout", + ( + "rowwise", + "gate_first", + "linear_scales", + "scales_unpermuted", + "swizzled_input_scales", + "alpha_rolled", + ), +) +def test_wrong_layouts_are_rejected(layout): + """Each operand prepared the wrong way -- trtllm-gen's row interleave, gate + rows first, scales left linear, scales not given the weight's row + permutation, the activation's scales swizzled, alpha off by one expert -- + returns finite, plausibly scaled output that `_compare` must reject against + the reference of the correctly prepared operands.""" + tokens, tile = 128, 128 + x, x_sf = _activation(tokens, seed=3) + sort = _sort(_ids("router", tokens, tile, seed=4), tile) + good = _args(x, x_sf, sort, tile, GSF["calibrated"]) + if layout == "swizzled_input_scales": + _xs, sf_swz = _activation(tokens, seed=3, swizzled=True) + bad = dict(good, input_scale=sf_swz.reshape(tokens, H // VEC)) + elif layout == "alpha_rolled": + bad = dict(good, alpha=good["alpha"].roll(1).contiguous()) + else: + bad = _args(x, x_sf, sort, tile, GSF["calibrated"], layout=layout) + out = op(**bad) + ref = op.reference(**good) + with pytest.raises(AssertionError): + _compare(out, ref) + + +# ── guards and loud rejections ──────────────────────────────────────────── + + +def _strided(t: torch.Tensor) -> torch.Tensor: + """The same values as `t`, as a non-contiguous view into a wider buffer.""" + u = t.view(torch.uint8) if t.element_size() == 1 else t + if u.dim() == 1: + wide = torch.zeros(2 * u.numel(), dtype=u.dtype, device=DEV) + wide[::2] = u + view = wide[::2] + else: + wide = torch.zeros(*u.shape[:-1], 2 * u.shape[-1], dtype=u.dtype, device=DEV) + wide[..., : u.shape[-1]] = u + view = wide[..., : u.shape[-1]] + return view.view(t.dtype) if t.element_size() == 1 else view + + +def _u8(t: torch.Tensor) -> torch.Tensor: + return t.view(torch.uint8) if t.element_size() == 1 else t + + +@pytest.mark.parametrize("name", ("input", "input_scale", "weight", "weight_scale", "alpha")) +def test_noncontiguous_operand_answers_wrongly_and_is_refused(name): + """The runner hands data_ptr() to the kernel and never reads a stride: the + same values through a strided view come back finite and wrong.""" + tokens, tile = 64, 128 + args, _ = _problem(tokens, tile, seed=21) + ref = op.reference(**args) + bad = dict(args, **{name: _strided(args[name])}) + assert not bad[name].is_contiguous() + assert torch.equal(_u8(bad[name]), _u8(args[name])) + with pytest.raises(AssertionError): + _compare(op(**bad), ref) + with validating(op), pytest.raises(AssertionError, match=f"^{name} must be contiguous"): + op(**bad) + + +def test_short_alpha_answers_wrongly_and_is_refused(): + """A 1-element alpha (one per layer instead of one per expert): the kernel reads + alpha[expert] past its end without raising -- here into zeros of the same + allocation, so every expert but the first computes zero.""" + tokens, tile = 64, 128 + args, _ = _problem(tokens, tile, seed=31) + ref = op.reference(**args) + buf = torch.zeros(E, dtype=torch.float32, device=DEV) + buf[0] = args["alpha"][0] + bad = dict(args, alpha=buf[:1]) + with pytest.raises(AssertionError): + _compare(op(**bad), ref) + with ( + validating(op), + pytest.raises(AssertionError, match="alpha has 1 elements for 512 experts"), + ): + op(**bad) + + +def test_short_weight_scale_answers_wrongly_and_is_refused(): + """A weight_scale holding half the bytes the weight needs: the kernel reads past + its end without raising -- here into zeros of the same allocation, so the upper + half of the experts compute zero. Past the end of an allocation the same call + faults with an illegal memory access.""" + tokens, tile = 64, 128 + args, _ = _problem(tokens, tile, seed=32) + ref = op.reference(**args) + ws = args["weight_scale"].reshape(-1) + half = ws.numel() // 2 + buf = torch.zeros(ws.numel(), dtype=torch.uint8, device=DEV) + buf[:half] = ws[:half] + bad = dict(args, weight_scale=buf[:half].view(E, INTER, H // VEC)) + with pytest.raises(AssertionError): + _compare(op(**bad), ref) + with validating(op), pytest.raises(AssertionError, match="weight_scale holds"): + op(**bad) + + +def test_loud_rejections(): + """Raised before the launch: a tile size other than the sort's, either way + round; 64, which the runner does not build; a wrong scale width, a 2-element + global_sf, a bf16 alpha, scales typed e4m3 and an activation typed uint8.""" + tokens = 16 + args, _ = _problem(tokens, 128, seed=41) + with pytest.raises(AssertionError): + op(**dict(args, tile_size=256)) + args256, _ = _problem(tokens, 256, seed=42) + with pytest.raises(AssertionError): + op(**dict(args256, tile_size=128)) + with pytest.raises(ValueError, match="Tile size 64 is not supported"): + op(**dict(args, tile_size=64)) + with pytest.raises(AssertionError): + op(**dict(args, input_scale=args["input_scale"][:, :-4].contiguous())) + with pytest.raises(AssertionError): + op(**dict(args, global_sf=torch.ones(2, dtype=torch.float32, device=DEV))) + with pytest.raises(AssertionError): + op(**dict(args, alpha=args["alpha"].bfloat16())) + with pytest.raises(AssertionError): + op(**dict(args, weight_scale=args["weight_scale"].view(torch.float8_e4m3fn))) + with pytest.raises(AssertionError): + op(**dict(args, input=args["input"].view(torch.uint8))) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_grouped_gemm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_grouped_gemm.py new file mode 100644 index 000000000000..00e319236451 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_grouped_gemm.py @@ -0,0 +1,733 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the CuteDSL NVFP4 FC2 entry of the Qwen3.8-Flash-Next target. + +The entry owns the contract. This file builds the checkpoint-shaped operands -- +512 experts' down_proj in the CuteDSL weight layout, the input as the CuteDSL +FC1 returns it (permuted NVFP4 rows, swizzled block scales, zero padding rows, +unwritten bytes past the live tiles), the sort as moe_sort returns it -- and +holds the claims no cell carries: per-pair invariance across tiles, tactics, +row order and prefixes; CUDA-graph replay of sort, FC1 and FC2 at the captured +row counts; the regions the op leaves unwritten; that the gate rejects every +wrong layout; the guards and the loud rejections. +""" + +import functools +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.cute_dsl_nvfp4_grouped_gemm import ( # noqa: E501 + Fc2Expected, + unswizzle, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.cute_dsl_nvfp4_grouped_gemm import ( # noqa: E501 + cute_dsl_nvfp4_grouped_gemm as op, +) +from tensorrt_llm._torch.autotuner import AutoTuner +from tensorrt_llm._torch.utils import ActivationType + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "the CuteDSL FC2 requires a CUDA device" + +DEV = "cuda" +E, K, H, INTER, VEC = 512, 10, 2560, 640, 16 +NB = INTER // VEC +#: The SwiGLU output maximum FC1's `global_sf` is calibrated to (see `_ckpt` in the entry). +ACT_AMAX = 48.0 +GSF = { + "calibrated": 448 * 6 / ACT_AMAX, + "saturating": 25 * 448 * 6 / ACT_AMAX, + "subnormal": 2.0**-12 * 448 * 6 / ACT_AMAX, +} +#: FC1's activation input scale in the chained tests: |x| = 5 maps to the e4m3 maximum. +ISC1 = 5.0 / (448 * 6) +SENTINEL = 0x5A5A + +#: Tie window of `_compare`, in units of the element's accumulation magnitude +#: (derived in its docstring). +TIE = 2.0**-20 +#: At most this fraction of the gated elements with a nonzero magnitude may take +#: the other side of a tie (4 of them when that is more). +FLIP_BUDGET = 1e-5 + +#: Row counts of a decode capture (requests x MTP verify width 1..4, up to 32 +#: requests), and the eager prefix lengths: both sides of the power-of-two edges up +#: to the 8192-row ceiling. +_CAPTURED_ROWS = (1, 2, 3, 4, 5, 8, 32, 96, 128) +_PREFIXES = (1, 2, 3, 4, 5, 8, 31, 32, 33, 127, 128, 129, 512, 1023, 1024, 1025, 4095, 4096, 8191) + + +# ── the gate ────────────────────────────────────────────────────────────── + + +def _bf16_order(x: torch.Tensor) -> torch.Tensor: + """bf16 -> an integer that is monotonic in its value, consecutive between + neighbouring bf16 values, and equal for +0 and -0.""" + bits = x.view(torch.int16).int() & 0xFFFF + mag = bits & 0x7FFF + return torch.where(bits >= 0x8000, -mag, mag) + + +def _deviation(out: torch.Tensor, ref: Fc2Expected, rows: torch.Tensor | None = None) -> dict: + """Every disagreement between the op and the reference on the rows it writes + (or on the subset `rows` of them). + + An element is `equal` when its value matches the reference's `out` (+0 and + -0 alike). Otherwise it is a `flip` when the two are neighbouring bf16 + values and the exact value lies within `TIE * magnitude` of the midpoint + between them, and `outside` when it is anything else. `far` counts the + disagreements that are not between neighbours at all. + """ + sel = ref.rows if rows is None else rows + assert not bool((sel & ~ref.rows).any()), "deviation reads written rows only" + span = ref.exact.shape[0] + idx = torch.nonzero(sel[:span]).squeeze(1) + o = out[idx] + r = ref.out[idx] + v = ref.exact[idx] + s = ref.magnitude[idx] + of, rf = o.double(), r.double() + finite = torch.isfinite(of) + diff = (of != rf) | ~finite + adj = diff & finite & ((_bf16_order(o) - _bf16_order(r)).abs() == 1) + dist = ((v - (of + rf) / 2).abs() / s.clamp_min(1e-300)).masked_fill(~adj, 0.0) + flip = adj & (dist <= TIE) + outside = diff & ~flip + return dict( + rows=int(idx.numel()), + elements=o.numel(), + nonzero=int((s > 0).sum()), + nonfinite=int((~finite).sum()), + flips=int(flip.sum()), + outside=int(outside.sum()), + far=int((diff & ~adj).sum()), + ) + + +def _compare(out: torch.Tensor, ref: Fc2Expected) -> None: + """Bitwise on every written row, except at bf16 rounding ties of derived width. + + The op accumulates the 640 products of a row in fp32 on the tensor cores -- + ten MMA k-steps of 64 -- multiplies by alpha in fp32 and rounds to bf16. Each + step rounds a partial sum no larger than the element's `magnitude` (`|alpha| * + sum |a_k w_k|`), so the op's fp32 value is within 11 roundings of `2^-24 * + magnitude` of the exact one; only where the exact value lies that close to a + bf16 midpoint may the op land on the other neighbour. Everywhere else the two + agree bit for bit. + + - Flip: only to the neighbouring bf16 value, and only within `TIE` = 2^-20 + (16 fp32 roundings) of `magnitude` from the midpoint. + - Flips: at most `FLIP_BUDGET` = 1e-5 of the gated elements whose magnitude is + nonzero (4 when that is fewer) -- padding rows do not widen the budget. + - No NaN or infinity. + + Every wrongly prepared operand puts most of the pair rows' elements outside the + gate, most of them not even a bf16 neighbour away. + """ + assert out.dtype == torch.bfloat16, out.dtype + assert out.shape == ref.out.shape, (tuple(out.shape), tuple(ref.out.shape)) + d = _deviation(out, ref) + assert d["rows"] > 0, d + assert d["nonfinite"] == 0, f"non-finite output on written rows: {d}" + assert d["outside"] == 0, f"elements off by more than a tie: {d}" + assert d["flips"] <= max(4, FLIP_BUDGET * d["nonzero"]), f"too many ties flipped: {d}" + + +# ── operands ────────────────────────────────────────────────────────────── + + +def _swizzle(sf: torch.Tensor) -> torch.Tensor: + """`[rows, cols]` scale bytes -> the flat 128x4-swizzled buffer; inverse of `unswizzle`.""" + rows, cols = sf.shape + return sf.reshape(rows // 128, 4, 32, cols // 4, 4).permute(0, 3, 2, 1, 4).reshape(-1) + + +def _quantized_experts(num_experts: int, rows: int, cols: int, seed: int): + """Per expert, a Gaussian `[rows, cols]` matrix quantized the way the checkpoint + stores it: NVFP4 codes, LINEAR e4m3 block scales, one `weight_scale_2`. Expert + magnitudes spread over 2^-1..2^1 so a wrong per-expert alpha shows.""" + g = torch.Generator(device=DEV).manual_seed(seed) + spread = torch.exp2(torch.rand(num_experts, generator=g, device=DEV) * 2 - 1) + w = torch.empty(num_experts, rows, cols // 2, dtype=torch.uint8, device=DEV) + s = torch.empty(num_experts, rows, cols // VEC, dtype=torch.uint8, device=DEV) + ws2 = torch.empty(num_experts, dtype=torch.float32, device=DEV) + for e in range(num_experts): + dense = torch.randn(rows, cols, generator=g, device=DEV) * (float(spread[e]) / cols**0.5) + dense = dense.to(torch.bfloat16) + gs = (448 * 6 / dense.abs().amax().float()).reshape(1) + q, sf = torch.ops.trtllm.fp4_quantize(dense, gs, VEC, False, False) + w[e] = q.view(rows, cols // 2) + s[e] = sf.view(rows, cols // VEC) + ws2[e] = 1.0 / gs + return w, s, ws2 + + +@functools.lru_cache(maxsize=None) +def _down_proj(): + """The checkpoint's `down_proj` experts: `[E, H, I/2]` codes, `[E, H, I/16]` + linear scales, `weight_scale_2` per expert.""" + return _quantized_experts(E, H, INTER, seed=20261013) + + +def _shuffle32(rows: int) -> torch.Tensor: + """trtllm-gen's 32-row block shuffle: destination row `8v + u` <- source `4u + v`.""" + idx = torch.arange(rows, device=DEV).view(rows // 32, 8, 4) + return idx.transpose(1, 2).reshape(-1) + + +@functools.lru_cache(maxsize=None) +def _weights(gsf: str = "calibrated", layout: str = "cutedsl"): + """`(weight, weight_scale, alpha)` in the layout CuteDSL reads -- what the target's + load-time transform produces -- or one of the wrong layouts the discrimination + test feeds it. `alpha = down_proj.input_scale * weight_scale_2`, with the + input scale the inverse of FC1's `global_sf`.""" + w, s, ws2 = _down_proj() + if layout == "shuffle32": + perm = _shuffle32(H) + w, s = w[:, perm], s[:, perm] + if layout == "linear_weight_scales": + s_k = s.contiguous() + else: + s_k = torch.ops.trtllm.block_scale_interleave(s.contiguous()).view(s.shape) + alpha = (ws2 / GSF[gsf]).contiguous() + return w.contiguous().view(torch.float4_e2m1fn_x2), s_k, alpha + + +def _pair_rows(pairs: int, seed: int, zero_pairs: int = 0) -> torch.Tensor: + """SwiGLU outputs of FC1, one bf16 `[I]` row per (token, slot) pair, in expanded + order: `up * silu(gate)` of Gaussian halves.""" + g = torch.Generator(device=DEV).manual_seed(seed) + up = torch.randn(pairs, INTER, generator=g, device=DEV) * 2 + gate = torch.randn(pairs, INTER, generator=g, device=DEV) * 2 + y = (up * gate * torch.sigmoid(gate)).to(torch.bfloat16) + y[:zero_pairs] = 0 + return y + + +def _fc2_input(y: torch.Tensor, sort, gsf: str, garbage_seed: int = 0): + """This op's input as the CuteDSL FC1 returns it: each pair's row NVFP4-quantized + with `global_sf` at its permuted row, block scales 128x4-swizzled, the padding + rows of live tiles zero, and the rows past the live tiles holding stale bytes.""" + t2e, _t2mn, e2p, p2e, _total, n_tiles = sort + m = p2e.numel() + tile = m // t2e.numel() + span = int(n_tiles) * tile + x = torch.zeros(m, INTER, dtype=torch.bfloat16, device=DEV) + x[e2p.reshape(-1).long()[: y.shape[0]]] = y + gs = torch.tensor([GSF[gsf]], dtype=torch.float32, device=DEV) + q, sf = torch.ops.trtllm.fp4_quantize(x, gs, VEC, False, False) + q = q.view(torch.uint8).reshape(m, INTER // 2) + sf = sf.view(torch.uint8).reshape(m, NB) + g = torch.Generator(device=DEV).manual_seed(garbage_seed) + q[span:] = torch.randint(0, 256, q[span:].shape, generator=g, device=DEV, dtype=torch.uint8) + sf[span:] = torch.randint(0, 256, sf[span:].shape, generator=g, device=DEV, dtype=torch.uint8) + return q.view(torch.float4_e2m1fn_x2), _swizzle(sf).contiguous() + + +def _ids(dist: str, tokens: int, tile: int, seed: int) -> torch.Tensor: + """Expert ids `[tokens, K]` int32, every row naming K distinct experts.""" + g = torch.Generator(device=DEV).manual_seed(seed) + if dist == "router": + scores = torch.rand(tokens, E, device=DEV, generator=g) + return scores.topk(K, dim=-1).indices.to(torch.int32) + if dist == "hot": + hot = torch.randperm(E, device=DEV, generator=g)[:K] + slots = torch.rand(tokens, K, device=DEV, generator=g).argsort(dim=-1) + return hot[slots].to(torch.int32) + if dist == "tile_edges": + edge = (7, 300, 511) + scores = torch.rand(tokens, E, device=DEV, generator=g) + scores[:, list(edge)] = -1.0 + ids = scores.topk(K, dim=-1).indices.to(torch.int32) + for slot, (expert, rows) in enumerate(zip(edge, (tile - 1, tile, tile + 1))): + ids[:rows, slot] = expert + return ids + if dist == "round_robin": + base = torch.arange(tokens, device=DEV).unsqueeze(1) * K + return ((base + torch.arange(K, device=DEV)) % E).to(torch.int32) + raise ValueError(dist) + + +def _sort(ids: torch.Tensor, tile: int): + """moe_sort as the CuteDSL MoE calls it; the routing weights do not reach FC2.""" + scales = torch.full(ids.shape, 0.1, dtype=torch.float32, device=DEV) + return torch.ops.trtllm.moe_sort( + token_selected_experts=ids, + token_final_scales=scales, + num_experts=E, + top_k=K, + local_expert_offset=0, + local_num_experts=E, + tile_tokens_dim=tile, + ) + + +def _args(x, x_sf, sort, tile, gsf="calibrated", layout="cutedsl"): + """The keyword arguments the target passes, by the op's own names.""" + weight, weight_scale, alpha = _weights(gsf, layout) + return dict( + input=x, + weight=weight, + input_scale=x_sf, + weight_scale=weight_scale, + alpha=alpha, + tile_idx_to_group_idx=sort[0], + num_non_exiting_tiles=sort[5], + num_experts=E, + top_k=K, + num_local_experts=E, + local_expert_offset=0, + tile_size=tile, + output_dtype=torch.bfloat16, + scaling_vector_size=VEC, + ) + + +def _problem(tokens, tile, seed, dist="router", gsf="calibrated", zero_tokens=0): + sort = _sort(_ids(dist, tokens, tile, seed + 1), tile) + y = _pair_rows(tokens * K, seed, zero_pairs=zero_tokens * K) + x, x_sf = _fc2_input(y, sort, gsf, garbage_seed=seed + 2) + return _args(x, x_sf, sort, tile, gsf), sort + + +def _call(args, mma_n, cluster_n): + """One op call with the autotuner tactic whose MMA tile is `(tile_size, mma_n)` + and cluster `(tile_size // 128, cluster_n)`, selected through the tuner's own + capture/replay so the call still goes through the op.""" + tuner = AutoTuner.get() + with tuner.capture() as cap: + op(**args) + tile = args["tile_size"] + want = ((tile, mma_n), (tile // 128, cluster_n)) + for config in cap: + ((_runner, tactic),) = config + if tactic == want: + with tuner.replay(config): + return op(**args) + raise AssertionError(f"no tactic {want}") + + +def _pairs(out, sort, tokens): + """Each (token, slot) pair's output row, in expanded order -- what the combine + reads, whatever row the sort gave it.""" + return out[sort[2].reshape(-1).long()[: tokens * K]] + + +def _pair_mask(sort, m, tokens): + mask = torch.zeros(m, dtype=torch.bool, device=DEV) + mask[sort[2].reshape(-1).long()[: tokens * K]] = True + return mask + + +def _clone(v): + """A copy that does not depend on copy kernels for the fp4 dtype.""" + if not isinstance(v, torch.Tensor): + return v + return v.view(torch.uint8).clone().view(v.dtype) if v.element_size() == 1 else v.clone() + + +def _bits(t: torch.Tensor) -> torch.Tensor: + return t.view(torch.int16) + + +def _assert_pairs_equal(a, b, what): + assert torch.equal(_bits(a), _bits(b)), f"{what}: rows differ" + + +def _cell_id(cell): + s = cell.spec + return ( + f"t{s['tile_size']}-n{s['mma_n']}-c{s['cluster_n']}-T{s['tokens']}-{s['dist']}" + f"-{s['gsf']}-z{s['zero_tokens']}" + ) + + +# ── the certified cells ─────────────────────────────────────────────────── + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[_cell_id(c) for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + s = cell.spec + assert (s["num_experts"], s["top_k"], s["hidden"], s["intermediate"]) == (E, K, H, INTER) + seed = zlib.crc32(cell.why.encode()) + args, sort = _problem(s["tokens"], s["tile_size"], seed, s["dist"], s["gsf"], s["zero_tokens"]) + before = {k: _clone(v) for k, v in args.items() if isinstance(v, torch.Tensor)} + out = _call(args, s["mma_n"], s["cluster_n"]) + _compare(out, op.reference(**args)) + for k, v in before.items(): + assert torch.equal(_u8(args[k]), _u8(v)), f"{k} was written" + if s["zero_tokens"]: + zero = sort[2][: s["zero_tokens"]].reshape(-1).long() + assert bool((out[zero] == 0).all()), "an all-zero token's rows are not zero" + + +# ── invariance: a pair's row depends on the pair alone ──────────────────── + + +@pytest.mark.parametrize("tokens", (1, 128, 8192)) +def test_tiles_and_tactics_agree_per_pair(tokens, certifying): # noqa: F811 + """Tile 128 / 256, MMA N 128 / 256, cluster N 1 / 2 and the tuner's fallback + all produce every pair's row bit for bit the same, so a cell certifies the + tactic the autotuner picks in serving as well as the one it ran.""" + certifying(op) + ids = _ids("router", tokens, 128, seed=tokens + 1) + y = _pair_rows(tokens * K, seed=tokens) + seen = {} + for tile in (128, 256): + sort = _sort(ids, tile) + x, x_sf = _fc2_input(y, sort, "calibrated", garbage_seed=tile) + args = _args(x, x_sf, sort, tile) + for mma_n in (128, 256): + for cluster_n in (1, 2): + seen[(tile, mma_n, cluster_n)] = _pairs(_call(args, mma_n, cluster_n), sort, tokens) + fallback = op(**args) + seen[(tile, "fallback")] = _pairs(fallback, sort, tokens) + _compare(fallback, op.reference(**args)) + first = seen[(128, 128, 1)] + for key, got in seen.items(): + _assert_pairs_equal(got, first, f"T={tokens} {key} vs (128, 128, 1)") + + +@pytest.mark.parametrize("tile", (128, 256)) +def test_row_order_and_repeat_do_not_change_a_pair(tile, certifying): # noqa: F811 + """moe_sort orders an expert's rows by atomics (from T = 5). Moving every pair to + another row of its expert's run, and calling twice, leaves each pair's row + bitwise unchanged: the FC2 result is reproducible run to run.""" + certifying(op) + tokens = 4096 + ids = _ids("router", tokens, tile, seed=78) + sort = _sort(ids, tile) + y = _pair_rows(tokens * K, seed=77) + x, x_sf = _fc2_input(y, sort, "calibrated") + args = _args(x, x_sf, sort, tile) + base = _pairs(op(**args), sort, tokens) + _assert_pairs_equal(_pairs(op(**args), sort, tokens), base, "repeat call") + + t2e, t2mn, e2p, p2e, n_tiles = sort[0], sort[1], sort[2], sort[3], int(sort[5]) + g = torch.Generator(device=DEV).manual_seed(5) + p2e_s = p2e.clone() + starts = torch.arange(n_tiles, device=DEV) * tile + run_of = torch.full((p2e.numel(),), -1, dtype=torch.long, device=DEV) + for i, (start, limit) in enumerate(zip(starts.tolist(), t2mn[:n_tiles].tolist())): + run_of[start:limit] = int(t2e[i]) + for e in torch.unique(run_of[run_of >= 0]).tolist(): + rows = torch.nonzero(run_of == e).squeeze(1) + p2e_s[rows] = p2e[rows[torch.randperm(rows.numel(), generator=g, device=DEV)]] + e2p_s = e2p.clone().reshape(-1) + real = run_of >= 0 + e2p_s[p2e_s[real].long()] = torch.nonzero(real).squeeze(1).to(torch.int32) + shuffled = (t2e, t2mn, e2p_s.view_as(e2p), p2e_s, sort[4], sort[5]) + assert not torch.equal(e2p_s.view_as(e2p), e2p), "the shuffle moved no pair" + xs, xs_sf = _fc2_input(y, shuffled, "calibrated") + args_s = _args(xs, xs_sf, shuffled, tile) + out_s = op(**args_s) + _compare(out_s, op.reference(**args_s)) + _assert_pairs_equal(_pairs(out_s, shuffled, tokens), base, "pairs moved within their runs") + + +@pytest.mark.parametrize("tile", (128, 256)) +def test_eager_rows_prefix_bitwise(tile, certifying): # noqa: F811 + """The 8192-row call is gated against the reference, and every prefix call's + pairs are bitwise the 8192-row call's.""" + certifying(op) + tokens = 8192 + ids = _ids("router", tokens, tile, seed=9) + y = _pair_rows(tokens * K, seed=8) + sort = _sort(ids, tile) + x, x_sf = _fc2_input(y, sort, "calibrated") + args = _args(x, x_sf, sort, tile) + full = op(**args) + _compare(full, op.reference(**args)) + full_pairs = _pairs(full, sort, tokens) + for t in _PREFIXES: + sub_sort = _sort(ids[:t].contiguous(), tile) + xs, xs_sf = _fc2_input(y[: t * K], sub_sort, "calibrated", garbage_seed=t) + got = _pairs(op(**_args(xs, xs_sf, sub_sort, tile)), sub_sort, t) + _assert_pairs_equal(got, full_pairs[: t * K], f"prefix T={t}") + + +# ── chained from the CuteDSL FC1, eager and captured ────────────────────── + + +@functools.lru_cache(maxsize=None) +def _fc1_weights(): + """FC1 operands for the chain: `[E, 2I, H/2]` codes, swizzled scales, per-expert + alpha and the layer's `global_sf`. Their layout is FC1's business; what FC2 + sees is FC1's output, whatever the weights.""" + w, s, ws2 = _quantized_experts(E, 2 * INTER, H, seed=20261014) + s_k = torch.ops.trtllm.block_scale_interleave(s.contiguous()).view(s.shape) + alpha = (ISC1 * ws2).contiguous() + gsf = torch.tensor([GSF["calibrated"]], dtype=torch.float32, device=DEV) + return w.view(torch.float4_e2m1fn_x2), s_k, alpha, gsf + + +def _activation(tokens: int, seed: int): + """The MoE input as the target quantizes it: `fp4_quantize(x, 1/isc1, 16, False, + False)`, scales viewed `[T, H/16]`.""" + g = torch.Generator(device=DEV).manual_seed(seed) + x = torch.randn(tokens, H, generator=g, device=DEV).to(torch.bfloat16) + gs = torch.tensor([1.0 / ISC1], dtype=torch.float32, device=DEV) + q, sf = torch.ops.trtllm.fp4_quantize(x, gs, VEC, False, False) + return q.view(torch.float4_e2m1fn_x2), sf.view(tokens, H // VEC) + + +def _fc1(xq, xsf, sort, tile): + """The CuteDSL FC1, called as the target calls it; returns `(c, c_sf)`.""" + w1, s1, alpha1, gsf = _fc1_weights() + return torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell( + input=xq, + weight=w1, + input_scale=xsf, + weight_scale=s1, + alpha=alpha1, + tile_idx_to_group_idx=sort[0], + tile_idx_to_mn_limit=sort[1], + permuted_idx_to_expanded_idx=sort[3], + num_non_exiting_tiles=sort[5], + global_sf=gsf, + num_experts=E, + top_k=K, + num_local_experts=E, + local_expert_offset=0, + tile_size=tile, + scaling_vector_size=VEC, + activation_type=int(ActivationType.Swiglu), + ) + + +@pytest.mark.parametrize("tile", (128, 256)) +def test_fc1_output_feeds_it_at_prefill(tile, certifying): # noqa: F811 + """The real FC1 output at 8192 tokens -- the target's prefill -- passes `_compare`, + and the padding rows FC1 zeroes come out zero.""" + certifying(op) + tokens = 8192 + xq, xsf = _activation(tokens, seed=91) + sort = _sort(_ids("router", tokens, tile, seed=92), tile) + h, h_sf = _fc1(xq, xsf, sort, tile) + args = _args(h, h_sf, sort, tile) + out = op(**args) + ref = op.reference(**args) + _compare(out, ref) + pad = ref.rows & ~_pair_mask(sort, out.shape[0], tokens) + assert int(pad.sum()) > 0 + assert bool((out[pad] == 0).all()), "a padding row of FC1's output did not come out zero" + + +def _capture(fn): + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for _ in range(2): + fn() + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + static = fn() + return graph, static + + +@pytest.mark.parametrize("tile", (128, 256)) +def test_captured_rows_replay_on_new_data(tile, certifying): # noqa: F811 + """moe_sort, FC1 and this op are captured together, as the target's graph holds + them; each graph is replayed on two fresh activations and id sets, and its + written rows are gated against the reference of the graph's own FC1 output and + sort, and bitwise against an eager call on them.""" + certifying(op) + for tokens in _CAPTURED_ROWS: + xq, xsf = _activation(tokens, seed=tokens) + ids = _ids("router", tokens, tile, seed=500 + tokens) + scales = torch.full(ids.shape, 0.1, dtype=torch.float32, device=DEV) + state = {} + + def step(xq=xq, xsf=xsf, ids=ids, scales=scales, state=state, tile=tile): + sort = torch.ops.trtllm.moe_sort( + token_selected_experts=ids, + token_final_scales=scales, + num_experts=E, + top_k=K, + local_expert_offset=0, + local_num_experts=E, + tile_tokens_dim=tile, + ) + h, h_sf = _fc1(xq, xsf, sort, tile) + state["sort"], state["h"], state["h_sf"] = sort, h, h_sf + return op(**_args(h, h_sf, sort, tile)) + + graph, static_out = _capture(step) + for trial in (1, 2): + nx, nx_sf = _activation(tokens, seed=1000 * trial + tokens) + xq.view(torch.uint8).copy_(nx.view(torch.uint8)) + xsf.copy_(nx_sf) + ids.copy_(_ids("router", tokens, tile, seed=7000 * trial + tokens)) + graph.replay() + args = _args(state["h"], state["h_sf"], state["sort"], tile) + ref = op.reference(**args) + _compare(static_out, ref) + eager = op(**{k: _clone(v) for k, v in args.items()}) + rows = ref.rows + assert torch.equal(_bits(static_out[rows]), _bits(eager[rows])), (tokens, trial) + del graph, static_out + torch.cuda.synchronize() + + +# ── what the op writes ──────────────────────────────────────────────────── + + +@pytest.mark.parametrize("tile", (128, 256)) +@pytest.mark.parametrize("tokens", (3, 300)) +def test_unwritten_regions(tokens, tile): + """A sentinel planted in a captured graph's output: which rows replay writes. + Every row of the live tiles is written -- the padding rows computed from + whatever the input holds there, zero for FC1's zeroed padding -- and the rows + past `num_non_exiting_tiles` are not: under a graph they keep whatever the + buffer held.""" + args, sort = _problem(tokens, tile, seed=11) + graph, out = _capture(lambda: op(**args)) + m = out.shape[0] + span = int(sort[5]) * tile + pairs = _pair_mask(sort, m, tokens) + pad = torch.zeros(m, dtype=torch.bool, device=DEV) + pad[:span] = ~pairs[:span] + assert int(pad.sum()) > 0 + # With no more pairs than experts the sort sizes one tile per pair; when no + # two pairs share an expert (T = 3 here) every tile is live. + assert span < m or tokens * K <= E, "the sort left no tile unused" + _bits(out).fill_(SENTINEL) + graph.replay() + _compare(out, op.reference(**args)) + untouched = (_bits(out) == SENTINEL).all(dim=1) + assert bool(untouched[span:].all()), "a row past num_non_exiting_tiles was written" + assert not bool(untouched[:span].any()), "a row of a live tile was not written" + assert bool((out[pad] == 0).all()), "a zero padding row did not come out zero" + pair_rows = out[pairs].clone() + + # Nonzero codes (all 0.5) under unit block scales in the padding rows. + codes = args["input"].view(torch.uint8) + codes[pad] = 0x11 + sf = unswizzle(args["input_scale"], m, NB).clone() + sf[pad] = 0x38 + args["input_scale"].copy_(_swizzle(sf)) + graph.replay() + assert bool((out[pad] != 0).any()), "nonzero padding input still gave zero rows" + assert torch.equal(_bits(out[pairs]), _bits(pair_rows)), "padding input changed a pair row" + _compare(out, op.reference(**args)) + del graph + + +# ── the gate sees a wrong layout ────────────────────────────────────────── + + +@pytest.mark.parametrize( + "layout", + ( + "shuffle32", + "linear_weight_scales", + "linear_input_scales", + "alpha_rolled", + "alpha_per_layer", + ), +) +def test_wrong_layouts_are_rejected(layout): + """Each operand prepared the wrong way -- trtllm-gen's 32-row shuffle, weight + scales left linear, the input's scales linear, alpha off by one expert, one + alpha for the whole layer -- returns finite, plausibly scaled output that + `_compare` must reject against the reference of the correctly prepared + operands.""" + tokens, tile = 128, 128 + good, sort = _problem(tokens, tile, seed=3) + m = good["input"].shape[0] + if layout in ("shuffle32", "linear_weight_scales"): + bad = _args(good["input"], good["input_scale"], sort, tile, layout=layout) + elif layout == "linear_input_scales": + linear = unswizzle(good["input_scale"], m, NB).reshape(-1).contiguous() + bad = dict(good, input_scale=linear) + elif layout == "alpha_rolled": + bad = dict(good, alpha=good["alpha"].roll(1).contiguous()) + else: + bad = dict(good, alpha=torch.full_like(good["alpha"], float(good["alpha"].mean()))) + out = op(**bad) + ref = op.reference(**good) + assert _deviation(out, ref, rows=_pair_mask(sort, m, tokens))["nonfinite"] == 0 + with pytest.raises(AssertionError): + _compare(out, ref) + + +# ── guards and loud rejections ──────────────────────────────────────────── + + +def _u8(t: torch.Tensor) -> torch.Tensor: + return t.view(torch.uint8) if t.element_size() == 1 else t + + +def _strided(t: torch.Tensor) -> torch.Tensor: + """The same values as `t`, as a non-contiguous view into a wider buffer.""" + u = _u8(t) + if u.dim() == 1: + wide = torch.zeros(2 * u.numel(), dtype=u.dtype, device=DEV) + wide[::2] = u + view = wide[::2] + else: + wide = torch.zeros(*u.shape[:-1], 2 * u.shape[-1], dtype=u.dtype, device=DEV) + wide[..., : u.shape[-1]] = u + view = wide[..., : u.shape[-1]] + return view.view(t.dtype) if t.element_size() == 1 else view + + +@pytest.mark.parametrize( + "name", + ("input", "input_scale", "weight", "weight_scale", "alpha", "tile_idx_to_group_idx"), +) +def test_noncontiguous_operand_answers_wrongly_and_is_refused(name): + """The runner hands data_ptr() to the kernel and never reads a stride: the + same values through a strided view come back finite and wrong.""" + tokens, tile = 64, 128 + args, _ = _problem(tokens, tile, seed=21) + ref = op.reference(**args) + bad = dict(args, **{name: _strided(args[name])}) + assert not bad[name].is_contiguous() + assert torch.equal(_u8(bad[name]), _u8(args[name])) + with pytest.raises(AssertionError): + _compare(op(**bad), ref) + with validating(op), pytest.raises(AssertionError, match=f"^{name} must be contiguous"): + op(**bad) + + +def test_loud_rejections(): + """Raised before the launch: a tile size other than the sort's, either way + round; 64, which the runner does not build; a non-bf16 output; and a wrong + dtype, rank or size of every tensor operand, each tripping a runner assert.""" + tokens = 16 + args, _ = _problem(tokens, 128, seed=41) + m = args["input"].shape[0] + with pytest.raises(AssertionError): + op(**dict(args, tile_size=256)) + args256, _ = _problem(tokens, 256, seed=42) + with pytest.raises(AssertionError): + op(**dict(args256, tile_size=128)) + with pytest.raises(ValueError, match="supports tile_size .* 128 and 256 only, but got 64"): + op(**dict(args, tile_size=64)) + rejected = dict( + output_dtype=torch.float16, + alpha=args["alpha"].bfloat16(), + input=args["input"].view(torch.uint8), + input_scale=args["input_scale"].view(m, NB), + weight_scale=args["weight_scale"].view(torch.float8_e4m3fn), + scaling_vector_size=32, + tile_idx_to_group_idx=args["tile_idx_to_group_idx"].long(), + num_non_exiting_tiles=torch.zeros(2, dtype=torch.int32, device=DEV), + ) + short = dict( + alpha=args["alpha"][:1].contiguous(), + weight_scale=args["weight_scale"][: E // 2], + input_scale=args["input_scale"][:-4].contiguous(), + tile_idx_to_group_idx=args["tile_idx_to_group_idx"][:-1].contiguous(), + ) + for name, value in list(rejected.items()) + list(short.items()): + with pytest.raises(AssertionError): + op(**dict(args, **{name: value})) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_weight_path.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_weight_path.py new file mode 100644 index 000000000000..89d85973e85f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_weight_path.py @@ -0,0 +1,362 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The routed experts of the Qwen3.8-Flash-Next target, from checkpoint-format tensors to output. + +The CuteDSL op tests build their operands with layout code of their own. This file builds +the checkpoint's expert tensors (512 experts: packed NVFP4 `up_proj` / `gate_proj` / +`down_proj` codes, linear e4m3 block scales, one `weight_scale_2` per expert, per-expert +`input_scale`s) and runs them through the target's own load-time transforms +(`weights._expert_fc1_w` / `_expert_fc1_s` / `_expert_fc2_w` / `_expert_fc2_s`) and scalar +derivation (`modeling._moe_layer_scalars`), once per routed backend. The operands then go +through the call sequence of the target's `_moe_block` -- CuteDSL: sort, FC1 with SwiGLU, +FC2, unpermute; trtllm-gen: the one-call runner -- with explicit routing. + +Two routings hit all 512 experts. In the first every token's only non-zero routing weight is +on its own expert, so each expert's output is checked alone: a layout error confined to one +expert (a permutation, a scale swizzle, an alpha) cannot hide in a sum over ten. The second is +a weighted top-10. Both backends are held to a float64 reference over the dequantized +checkpoint tensors, with `down_proj.input_scale` both spread across the experts and uniform. +When it is uniform the two backends requantize the SwiGLU output identically, so they are +also held to each other. + +The reference does not model the kernels' accumulation. It models the one rounding step +that is not exact: the SwiGLU output is requantized to NVFP4 before FC2, at the backend's +own global scale (the layer's single `global_sf` on CuteDSL, each expert's +`1 / down_proj.input_scale` on trtllm-gen). +""" + +import functools +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fp4_block_scale_moe_runner import ( + quantize_nvfp4, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 + modeling as M, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 + weights as W, +) + +assert torch.cuda.is_available(), "the CuteDSL experts require a CUDA device" + +DEV = "cuda" +E, K, H, INTER, VEC = 512, 10, 2560, 640, 16 +#: The activation's input scale, one number for all experts: |x| = 5 maps to the e4m3 maximum. +ISC1 = 5.0 / (448 * 6) +#: The SwiGLU output maximum `down_proj.input_scale` is calibrated to. With `spread` each +#: expert's own value ranges over 2^-1/2 .. 2^1/2 of it, so the layer's `max` is not every +#: expert's. +ISC2 = 64.0 / (448 * 6) +E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) + +#: Tokens of the two routings: one per expert, and a weighted top-10. +ONE_HOT_TOKENS = E +WEIGHTED_TOKENS = 1024 + +BACKENDS = ("CUTEDSL", "TRTLLM") +CASES = ("one_hot", "weighted") +#: The float64 reference's gate, relative per token. Both backends sit at 2.5e-3 or below: +#: the bf16 rounding of the output is 1e-3 of it, the rest is the reference's fp32 / tie +#: differences. A wrong layout is ~1; a wrong per-expert scalar is its own relative error. +REFERENCE_TOLERANCE = 5e-3 +#: The same, between the backends under a uniform `down_proj.input_scale`. +BACKEND_TOLERANCE = 3e-3 + + +# ── the checkpoint ──────────────────────────────────────────────────────── + + +def _quantize(dense: torch.Tensor): + """A bf16 `[R, C]` matrix as the checkpoint stores it: packed e2m1 codes `[R, C/2]` + uint8, e4m3 block scales `[R, C/16]`, and the `weight_scale_2` that undoes the global + scale the block scales were derived under.""" + rows, cols = dense.shape + gs = (448 * 6 / dense.abs().amax().float()).reshape(1) + codes, sf = torch.ops.trtllm.fp4_quantize(dense, gs, VEC, False, False) + return ( + codes.view(torch.uint8).view(rows, cols // 2), + sf.view(torch.uint8).view(rows, cols // VEC).view(torch.float8_e4m3fn), + (1.0 / gs).squeeze(0), + ) + + +@functools.lru_cache(maxsize=None) +def _checkpoint(spread: bool) -> SimpleNamespace: + """All 512 experts of one layer, with the checkpoint's tensor names and dtypes. + + `up_proj` and `gate_proj` are quantized together, so they share one `weight_scale_2`, as + the fused FC1 requires. Expert magnitudes spread over 2^-1 .. 2^1/2 (FC1) and 2^-1 .. 2^1 + (FC2), so a wrong per-expert scalar shows.""" + g = torch.Generator(device=DEV).manual_seed(20261014) + spread1 = torch.exp2(torch.rand(E, generator=g, device=DEV) * 1.5 - 1) + spread2 = torch.exp2(torch.rand(E, generator=g, device=DEV) * 2 - 1) + fields = {k: [] for k in ("up_w", "up_s", "gate_w", "gate_s", "down_w", "down_s")} + ws2_gu, ws2_dn = [], [] + for e in range(E): + gate_up = torch.randn(2 * INTER, H, generator=g, device=DEV) * (spread1[e] / H**0.5) + down = torch.randn(H, INTER, generator=g, device=DEV) * (spread2[e] / INTER**0.5) + codes, sf, ws2 = _quantize(gate_up.to(torch.bfloat16)) + fields["up_w"].append(codes[:INTER]) + fields["gate_w"].append(codes[INTER:]) + fields["up_s"].append(sf[:INTER]) + fields["gate_s"].append(sf[INTER:]) + ws2_gu.append(ws2) + codes, sf, ws2 = _quantize(down.to(torch.bfloat16)) + fields["down_w"].append(codes) + fields["down_s"].append(sf) + ws2_dn.append(ws2) + ws2_gu, ws2_dn = torch.stack(ws2_gu), torch.stack(ws2_dn) + isc1 = torch.full((E,), ISC1, dtype=torch.float32, device=DEV) + isc2 = torch.full((E,), ISC2, dtype=torch.float32, device=DEV) + if spread: + isc2 = isc2 * torch.exp2(torch.rand(E, generator=g, device=DEV) - 0.5) + return SimpleNamespace( + **fields, + ws2_gu=ws2_gu, + ws2_dn=ws2_dn, + isc2=isc2, + # The rows `modeling._moe_layer_scalars` reads, under layer 0. + scalars={ + "l0_e_isc1": isc1, + "l0_e_isc1_up": isc1.clone(), + "l0_e_ws2_1": ws2_gu, + "l0_e_ws2_1_up": ws2_gu.clone(), + "l0_e_isc2": isc2, + "l0_e_ws2_2": ws2_dn, + }, + ) + + +# ── what the target builds from it ──────────────────────────────────────── + + +@functools.lru_cache(maxsize=None) +def _layer(backend: str, spread: bool) -> SimpleNamespace: + """The layer's expert operands for `backend`, built as the weight table builds them: + one transform call per expert and tensor, stacked, then viewed as `_layer_views` + views them.""" + ck = _checkpoint(spread) + core = SimpleNamespace(trunk_moe_backend=backend) + + def stack(fn, *tensors): + return torch.stack([fn(tuple(t[e] for t in tensors), core) for e in range(E)]) + + def stack1(fn, tensor): + return torch.stack([fn(tensor[e], core) for e in range(E)]) + + fc1_w = stack(W._expert_fc1_w, ck.up_w, ck.gate_w) + fc1_s = stack(W._expert_fc1_s, ck.up_s, ck.gate_s) + fc2_w = stack1(W._expert_fc2_w, ck.down_w) + fc2_s = stack1(W._expert_fc2_s, ck.down_s) + if backend == "CUTEDSL": + w_dtype, s_dtype = torch.float4_e2m1fn_x2, torch.uint8 + else: + w_dtype, s_dtype = torch.uint8, torch.float8_e4m3fn + return SimpleNamespace( + fc1_w=fc1_w.view(w_dtype), + fc1_s=fc1_s.view(s_dtype), + fc2_w=fc2_w.view(w_dtype), + fc2_s=fc2_s.view(s_dtype), + scalars=M._moe_layer_scalars(ck.scalars, 0, backend), + ) + + +def _cutedsl_routed(layer, x, ids, weights) -> torch.Tensor: + """The CuteDSL branch of `_moe_block`, from the quantization to the unpermute.""" + sc = layer.scalars + tokens = x.shape[0] + xq, xsf = M.fp4_quantize(x, sc["global_scale"], M._SF_VEC, False, M._SF_LINEAR) + tile_expert, tile_limit, expanded_to_permuted, permuted_to_expanded, _, live_tiles = M.moe_sort( + token_selected_experts=ids, + token_final_scales=weights, + num_experts=E, + top_k=K, + local_expert_offset=0, + local_num_experts=E, + tile_tokens_dim=M._CUTEDSL_TILE_SIZE, + ) + h, h_sf = M.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion( + input=xq.view(torch.float4_e2m1fn_x2), + weight=layer.fc1_w, + input_scale=xsf.view(tokens, H // M._SF_VEC), + weight_scale=layer.fc1_s, + alpha=sc["alpha1"], + tile_idx_to_group_idx=tile_expert, + tile_idx_to_mn_limit=tile_limit, + permuted_idx_to_expanded_idx=permuted_to_expanded, + num_non_exiting_tiles=live_tiles, + global_sf=sc["global_sf"], + num_experts=E, + top_k=K, + num_local_experts=E, + local_expert_offset=0, + tile_size=M._CUTEDSL_TILE_SIZE, + scaling_vector_size=M._SF_VEC, + activation_type=int(M.ActivationType.Swiglu), + ) + y = M.cute_dsl_nvfp4_grouped_gemm( + input=h, + weight=layer.fc2_w, + input_scale=h_sf, + weight_scale=layer.fc2_s, + alpha=sc["alpha2"], + tile_idx_to_group_idx=tile_expert, + num_non_exiting_tiles=live_tiles, + num_experts=E, + top_k=K, + num_local_experts=E, + local_expert_offset=0, + tile_size=M._CUTEDSL_TILE_SIZE, + output_dtype=torch.bfloat16, + scaling_vector_size=M._SF_VEC, + ) + return M.moe_unpermute_inplace( + permuted_input=y, + output=torch.empty(tokens, H, dtype=torch.bfloat16, device=x.device), + expanded_idx_to_permuted_idx=expanded_to_permuted, + topk_scales=weights, + ) + + +def _trtllm_gen_routed(layer, x, ids, weights) -> torch.Tensor: + """The trtllm-gen branch of `_moe_block`.""" + sc = layer.scalars + xq, xsf = M.fp4_quantize(x, sc["global_scale"], M._SF_VEC, False, M._SF_LINEAR) + return M.fp4_block_scale_moe_runner( + routing_logits=None, + routing_bias=None, + hidden_states=xq, + hidden_states_scale=xsf.view(torch.float8_e4m3fn), + gemm1_weights=layer.fc1_w, + gemm1_weights_scale=layer.fc1_s, + gemm1_bias=None, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=None, + gemm2_weights=layer.fc2_w, + gemm2_weights_scale=layer.fc2_s, + gemm2_bias=None, + output1_scale_scalar=sc["output1_scale_scalar"], + output1_scale_gate_scalar=sc["output1_scale_gate_scalar"], + output2_scale_scalar=sc["output2_scale_scalar"], + num_experts=E, + top_k=K, + n_group=None, + topk_group=None, + intermediate_size=INTER, + local_expert_offset=0, + local_num_experts=E, + routed_scaling_factor=None, + routing_method_type=M._ROUTING_METHOD_RENORM, + do_finalize=True, + act_type=M._ACT_TYPE_SWIGLU, + topk_weights=weights.to(torch.bfloat16), + topk_ids=ids, + )[0] + + +# ── the inputs and the reference ────────────────────────────────────────── + + +@functools.lru_cache(maxsize=None) +def _problem(case: str): + """`(x, ids, weights)`: bf16 activations, expert ids `[T, K]` int32 and routing weights + `[T, K]` fp32 (exact in bf16, which is the dtype trtllm-gen takes them in).""" + tokens = ONE_HOT_TOKENS if case == "one_hot" else WEIGHTED_TOKENS + g = torch.Generator(device=DEV).manual_seed(7 if case == "one_hot" else 8) + x = torch.randn(tokens, H, generator=g, device=DEV).to(torch.bfloat16) + scores = torch.rand(tokens, E, generator=g, device=DEV) + if case == "one_hot": + # The forced expert has the top score, so it sits in slot 0. + scores[torch.arange(tokens), torch.arange(tokens) % E] = 2.0 + weights = torch.zeros(tokens, K, device=DEV) + weights[:, 0] = 1.0 + else: + weights = torch.softmax(torch.randn(tokens, K, generator=g, device=DEV), dim=-1) + weights = weights.to(torch.bfloat16).float() + ids = scores.topk(K, dim=-1).indices.to(torch.int32) + return x, ids, weights + + +def _dequant(codes: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + """`[R, C/2]` packed e2m1 (element 2j in the low nibble) and `[R, C/16]` e4m3 scales + -> float64 `[R, C]`.""" + magnitude = torch.tensor(E2M1, dtype=torch.float64, device=codes.device) + table = torch.cat([magnitude, -magnitude]) + lo, hi = table[(codes & 0xF).long()], table[(codes >> 4).long()] + values = torch.stack([lo, hi], dim=-1).reshape(codes.shape[0], -1) + return values * scales.view(torch.float8_e4m3fn).double().repeat_interleave(VEC, dim=1) + + +def _reference(backend: str, case: str, spread: bool) -> torch.Tensor: + """The routed output in float64, `[T, H]`. + + The activation is what the kernels see, the dequantized NVFP4 of the target's own + quantization. Each expert's weights are dequantized from the checkpoint tensors, times + its `weight_scale_2`. The SwiGLU output is requantized at `backend`'s global scale.""" + ck, layer = _checkpoint(spread), _layer(backend, spread) + x, ids, weights = _problem(case) + sc = layer.scalars + xq, xsf = M.fp4_quantize(x, sc["global_scale"], M._SF_VEC, False, M._SF_LINEAR) + x_dq = _dequant(xq.view(torch.uint8), xsf.view(x.shape[0], H // VEC)) * ISC1 + if backend == "CUTEDSL": + g2 = [float(sc["global_sf"])] * E + else: + g2 = (1.0 / ck.isc2).tolist() + out = torch.zeros(x.shape[0], H, dtype=torch.float64, device=DEV) + for e in range(E): + token, slot = (ids == e).nonzero(as_tuple=True) + w = weights[token, slot].double() + token, w = token[w != 0], w[w != 0] + if token.numel() == 0: + continue + ws2 = ck.ws2_gu[e].double() + xe = x_dq[token] + up = xe @ (_dequant(ck.up_w[e], ck.up_s[e]) * ws2).T + gate = xe @ (_dequant(ck.gate_w[e], ck.gate_s[e]) * ws2).T + act = up * gate * torch.sigmoid(gate) + act = quantize_nvfp4(act.float(), g2[e]).double() / g2[e] + y = act @ (_dequant(ck.down_w[e], ck.down_s[e]) * ck.ws2_dn[e].double()).T + out.index_add_(0, token, y * w[:, None]) + return out + + +@functools.lru_cache(maxsize=None) +def _routed(backend: str, case: str, spread: bool) -> torch.Tensor: + x, ids, weights = _problem(case) + run = _cutedsl_routed if backend == "CUTEDSL" else _trtllm_gen_routed + return run(_layer(backend, spread), x, ids, weights).double() + + +def _token_error(out: torch.Tensor, ref: torch.Tensor) -> torch.Tensor: + """Each token's `|out - ref| / |ref|` (L2 over the hidden width).""" + return (out - ref).norm(dim=1) / ref.norm(dim=1).clamp_min(1e-12) + + +# ── tests ───────────────────────────────────────────────────────────────── + + +def test_routings_hit_every_expert() -> None: + for case in CASES: + _, ids, weights = _problem(case) + assert torch.unique(ids[weights != 0]).numel() == E, case + + +@pytest.mark.parametrize("spread", (True, False), ids=["isc2_spread", "isc2_uniform"]) +@pytest.mark.parametrize("case", CASES) +@pytest.mark.parametrize("backend", BACKENDS) +def test_routed_output_matches_the_float64_reference(backend: str, case: str, spread: bool) -> None: + err = _token_error(_routed(backend, case, spread), _reference(backend, case, spread)) + print(f"[err] {backend} {case} spread={spread}: mean {err.mean():.4e} max {err.max():.4e}") + assert err.max() <= REFERENCE_TOLERANCE, (err.mean().item(), err.max().item()) + + +@pytest.mark.parametrize("case", CASES) +def test_cutedsl_matches_trtllm_gen_under_a_uniform_input_scale(case: str) -> None: + err = _token_error(_routed("CUTEDSL", case, False), _routed("TRTLLM", case, False)) + print(f"[err] CUTEDSL vs TRTLLM {case}: mean {err.mean():.4e} max {err.max():.4e}") + assert err.max() <= BACKEND_TOLERANCE, (err.mean().item(), err.max().item()) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_sort.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_sort.py new file mode 100644 index 000000000000..9ba068a29ae8 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_sort.py @@ -0,0 +1,432 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the moe_sort entry of the Qwen3.8-Flash-Next target. + +The entry owns the contract. This file turns a cell into expert ids and holds +the claims no cell carries: CUDA-graph replay at the captured row counts, the +eager row sweep, which regions the op leaves unwritten, the token-order layout +of the block kernel, the scales being inert, the data hazards, the loud +rejections 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.moe_sort import ( # noqa: E501 + moe_sort as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "moe_sort requires a CUDA device" + +DEV = "cuda" +E, K = 512, 10 +SENTINEL = 0x5A5A5A5A + +_NAMES = ( + "tile_idx_to_expert_idx", + "tile_idx_to_mn_limit", + "expanded_idx_to_permuted_idx", + "permuted_idx_to_expanded_idx", + "total_num_padded_tokens", + "num_non_exiting_tiles", +) + + +def _compare(out: list[torch.Tensor], ref: list[torch.Tensor]) -> None: + """Integer outputs, so exact -- up to the one freedom the op has. + + Counts and the tile arrays' written prefix are bitwise. The order of an + expert's rows is the kernel's atomic order (it varies between identical calls + from T = 5), so `permuted_idx_to_expanded_idx` is gated per expert run as a + set: each real row's value sorted within its run must equal the reference's + run. `expanded_idx_to_permuted_idx` must then be the exact inverse: -1 + exactly where the expert is not local, and otherwise a real row of the right + run that maps back. Padding rows and tiles past `num_non_exiting_tiles` are + not written by the op and are not compared. + """ + assert len(out) == len(ref) == len(_NAMES), (len(out), len(ref)) + for name, o, r in zip(_NAMES, out, ref): + assert o.dtype == r.dtype == torch.int32, (name, o.dtype, r.dtype) + assert o.shape == r.shape, (name, tuple(o.shape), tuple(r.shape)) + t2e, t2mn, e2p, p2e, total, n_tiles = out + r_t2e, r_t2mn, r_e2p, r_p2e, r_total, r_n_tiles = ref + assert torch.equal(n_tiles.cpu(), r_n_tiles.cpu()), (n_tiles, r_n_tiles) + assert torch.equal(total.cpu(), r_total.cpu()), (total, r_total) + n = int(r_n_tiles) + assert torch.equal(t2e[:n], r_t2e[:n]), "tile_idx_to_expert_idx differs" + assert torch.equal(t2mn[:n], r_t2mn[:n]), "tile_idx_to_mn_limit differs" + + tile = p2e.numel() // t2e.numel() + num_expanded = e2p.numel() + real = r_p2e >= 0 + rows = torch.nonzero(real).squeeze(1) + run = r_t2e[rows // tile].long() + got = p2e[real].long() + assert bool(((got >= 0) & (got < num_expanded)).all()), ( + "permuted_idx_to_expanded_idx holds an out-of-range value on a real row" + ) + assert torch.equal( + torch.sort(run * num_expanded + got).values, run * num_expanded + r_p2e[real].long() + ), "an expert's run of permuted_idx_to_expanded_idx holds the wrong pairs" + + flat, r_flat = e2p.reshape(-1).long(), r_e2p.reshape(-1).long() + not_local = r_flat < 0 + assert bool((flat[not_local] == -1).all()), "a non-local pair is not -1" + p = flat[~not_local] + assert bool(((p >= 0) & (p < p2e.numel())).all()), "a local pair points outside the rows" + assert bool(real[p].all()), "a local pair points at a padding row" + assert torch.equal(p2e[p].long(), torch.arange(num_expanded, device=p.device)[~not_local]), ( + "expanded_idx_to_permuted_idx is not the inverse of permuted_idx_to_expanded_idx" + ) + + +def _ids(dist, tokens, tile, seed, num_experts=E, top_k=K): + """Expert ids, `[tokens, top_k]` int32, every row naming distinct experts.""" + g = torch.Generator(device=DEV).manual_seed(seed) + if dist == "router": + scores = torch.rand(tokens, num_experts, device=DEV, generator=g) + return scores.topk(top_k, dim=-1).indices.to(torch.int32) + if dist == "hot": + hot = torch.randperm(num_experts, device=DEV, generator=g)[:top_k] + slots = torch.rand(tokens, top_k, device=DEV, generator=g).argsort(dim=-1) + return hot[slots].to(torch.int32) + if dist == "tile_edges": + edge = (7, 300, 511) + scores = torch.rand(tokens, num_experts, device=DEV, generator=g) + scores[:, list(edge)] = -1.0 + ids = scores.topk(top_k, dim=-1).indices.to(torch.int32) + for slot, (expert, rows) in enumerate(zip(edge, (tile - 1, tile, tile + 1))): + ids[:rows, slot] = expert + return ids + if dist == "round_robin": + base = torch.arange(tokens, device=DEV).unsqueeze(1) * top_k + return ((base + torch.arange(top_k, device=DEV)) % num_experts).to(torch.int32) + if dist == "strided": + h = torch.randint(0, num_experts, (tokens, 1), device=DEV, generator=g) + return ((h + torch.arange(top_k, device=DEV) * 51) % num_experts).to(torch.int32) + raise ValueError(dist) + + +def _args(ids, tile, seed=0, scales_dtype=torch.float32): + """The CuteDSL call site's keyword arguments at the checkpoint's values.""" + g = torch.Generator(device=DEV).manual_seed(seed + 1) + scales = torch.rand(ids.shape, device=DEV, generator=g).to(scales_dtype) + return dict( + token_selected_experts=ids, + token_final_scales=scales, + num_experts=E, + top_k=K, + local_expert_offset=0, + local_num_experts=E, + tile_tokens_dim=tile, + ) + + +def _check(args): + """Drive the op and gate it through the entry's own reference and `_compare`.""" + out = op(**args) + ref = op.reference(**args) + _compare(out, ref) + return out, ref + + +def _cell_id(cell): + s = cell.spec + return f"t{s['tile_tokens_dim']}-T{s['tokens']}-{s['dist']}" + + +# ── the certified cells ─────────────────────────────────────────────────── + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[_cell_id(c) for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + s = cell.spec + seed = zlib.crc32(cell.why.encode()) + ids = _ids(s["dist"], s["tokens"], s["tile_tokens_dim"], seed) + args = _args(ids, s["tile_tokens_dim"], seed, s["scales_dtype"]) + assert (s["num_experts"], s["top_k"]) == (args["num_experts"], args["top_k"]) + assert (s["local_expert_offset"], s["local_num_experts"]) == (0, E) + ids_before = ids.clone() + scales_before = args["token_final_scales"].clone() + out, ref = _check(args) + assert torch.equal(ids, ids_before), "token_selected_experts was written" + assert torch.equal(args["token_final_scales"], scales_before), "token_final_scales was written" + if s["dist"] == "round_robin": + # every pair on its own expert: the shape bound is reached exactly + assert int(out[5]) == out[0].numel() == s["tokens"] * s["top_k"] + if s["dist"] == "tile_edges": + tile = s["tile_tokens_dim"] + n = int(ref[5]) + limits = out[1][:n].long() + starts = torch.arange(n, device=DEV) * tile + fill = (limits - starts).tolist() + experts = out[0][:n].tolist() + per_expert = {e: [] for e in (7, 300, 511)} + for e, f in zip(experts, fill): + if e in per_expert: + per_expert[e].append(f) + assert per_expert == {7: [tile - 1], 300: [tile], 511: [tile, 1]}, per_expert + + +# ── CUDA graphs: captured row counts, on new ids ────────────────────────── + +#: Row counts of a decode capture (requests x MTP verify width 1..4, up to 32 +#: requests): the block kernel (T <= 4), the first cluster-kernel T, and the widths +#: that follow. +_CAPTURED_ROWS = (1, 2, 3, 4, 5, 8, 32, 96, 128) + + +def _capture(args): + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for _ in range(2): + op(**args) + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + static_out = op(**args) + return graph, static_out + + +@pytest.mark.parametrize("tile", (128, 256)) +def test_captured_rows_replay_on_new_ids(tile, certifying): # noqa: F811 + """Each graph is replayed on two fresh id sets and gated against the reference, + and against an eager call on the same ids: counts and tile arrays bitwise, the + whole output bitwise where the kernel is deterministic (T <= 4).""" + certifying(op) + for rows in _CAPTURED_ROWS: + args = _args(_ids("router", rows, tile, seed=rows), tile, seed=rows) + graph, static_out = _capture(args) + for trial in (1, 2): + fresh = _ids("router", rows, tile, seed=1000 * trial + rows) + args["token_selected_experts"].copy_(fresh) + graph.replay() + eager = op(**dict(args, token_selected_experts=fresh.clone())) + ref = op.reference(**args) + _compare(static_out, ref) + _compare(eager, ref) + n = int(ref[5]) + for i in (0, 1): + assert torch.equal(static_out[i][:n], eager[i][:n]), (rows, trial, i) + for i in (4, 5): + assert torch.equal(static_out[i], eager[i]), (rows, trial, i) + if rows <= 4: + real = ref[3] >= 0 + assert torch.equal(static_out[2], eager[2]), (rows, trial) + assert torch.equal(static_out[3][real], eager[3][real]), (rows, trial) + del graph, static_out + torch.cuda.synchronize() + + +@pytest.mark.parametrize("tile", (128, 256)) +@pytest.mark.parametrize("rows", (3, 300)) +def test_padding_rows_and_spare_tiles_are_not_written(rows, tile): + """A sentinel planted in a captured graph's outputs survives replay in exactly + the regions `_compare` excludes; under a graph those regions keep the previous + replay's indices.""" + args = _args(_ids("router", rows, tile, seed=5), tile, seed=5) + graph, static_out = _capture(args) + for trial in (1, 2): + args["token_selected_experts"].copy_(_ids("router", rows, tile, seed=50 + trial)) + for t in static_out[:4]: + t.fill_(SENTINEL) + graph.replay() + ref = op.reference(**args) + _compare(static_out, ref) + n = int(ref[5]) + real = ref[3] >= 0 + assert bool((static_out[0][n:] == SENTINEL).all()) and bool( + (static_out[1][n:] == SENTINEL).all() + ) + assert bool((static_out[3][~real] == SENTINEL).all()) + assert not bool((static_out[3][real] == SENTINEL).any()) + assert not bool((static_out[2] == SENTINEL).any()), ( + "expanded_idx_to_permuted_idx is fully written" + ) + # replay on other ids without re-planting: the unwritten rows keep the last replay's values + before = static_out[3].clone() + args["token_selected_experts"].copy_(_ids("router", rows, tile, seed=99)) + graph.replay() + real = op.reference(**args)[3] >= 0 + assert torch.equal(static_out[3][~real], before[~real]) + if rows == 300: + stale = before[~real] + assert int(((stale >= 0) & (stale < rows * K)).sum()) > 0, "no valid-looking stale row" + del graph, static_out + + +# ── eager: the row sweep ────────────────────────────────────────────────── + +#: Both sides of every kernel-path edge: block / cluster at 4 / 5, the cluster +#: kernel's per-block work edges at 102 / 103 and 204 / 205 rows, its fast / slow +#: path edges at 3276 / 3277 and 6553 / 6554, and the 8192-row ceiling. +_SWEEP = ( + *(1, 2, 3, 4, 5, 6, 8, 31, 32, 33, 102, 103, 204, 205, 512, 1024), + *(3276, 3277, 4096, 6553, 6554, 8191, 8192), +) + + +@pytest.mark.parametrize("tile", (128, 256)) +def test_eager_row_sweep(tile, certifying): # noqa: F811 + certifying(op) + for tokens in _SWEEP: + _check(_args(_ids("router", tokens, tile, seed=tokens), tile, seed=tokens)) + + +# ── layout facts ────────────────────────────────────────────────────────── + + +def test_block_kernel_rows_follow_token_order(): + """T <= 4: each expert's rows in token order, bitwise the reference's.""" + for tile in (128, 256): + for tokens in (1, 2, 3, 4): + for seed in range(4): + args = _args(_ids("router", tokens, tile, seed=seed), tile, seed=seed) + out, ref = _check(args) + real = ref[3] >= 0 + assert torch.equal(out[2], ref[2]) and torch.equal(out[3][real], ref[3][real]) + + +def test_scales_are_inert_and_untouched(): + """The scales' values and dtype never reach an output, and are not written.""" + for tokens in (3, 300): + ids = _ids("router", tokens, 128, seed=21) + outs = [] + for scales in ( + torch.rand(tokens, K, device=DEV), + torch.rand(tokens, K, device=DEV).to(torch.bfloat16), + torch.zeros(tokens, K, device=DEV), + torch.full((tokens, K), float("nan"), device=DEV), + ): + bits = scales.view( + torch.int32 if scales.dtype == torch.float32 else torch.int16 + ).clone() + args = dict(_args(ids, 128), token_final_scales=scales) + out, ref = _check(args) + assert torch.equal(scales.view(bits.dtype), bits), "token_final_scales was written" + outs.append(out) + if tokens <= 4: # deterministic: every written element bitwise across the four + n, real = int(ref[5]), ref[3] >= 0 + for out in outs[1:]: + assert torch.equal(out[0][:n], outs[0][0][:n]) and torch.equal( + out[1][:n], outs[0][1][:n] + ) + assert torch.equal(out[2], outs[0][2]) and torch.equal( + out[3][real], outs[0][3][real] + ) + assert torch.equal(out[4], outs[0][4]) and torch.equal(out[5], outs[0][5]) + + +# ── data hazards is_valid cannot see ────────────────────────────────────── + + +def test_duplicate_expert_in_a_row(): + """Block kernel (T <= 4) keeps one of two slots naming the same expert; + the other paths keep both.""" + ids = _ids("router", 2, 128, seed=31) + ids[:, 1] = ids[:, 0] + args = _args(ids, 128) + out = op(**args) + ref = op.reference(**args) + for row in range(2): + pair = out[2][row, :2].tolist() + assert pair.count(-1) == 1, pair + assert int((out[2] >= 0).sum()) == 2 * K - 2 and int((ref[2] >= 0).sum()) == 2 * K + with pytest.raises(AssertionError): + _compare(out, ref) + + ids = _ids("router", 300, 128, seed=32) + ids[:, 1] = ids[:, 0] + _check(_args(ids, 128)) + + +def test_ids_past_int16_alias_off_the_block_kernel(): + """An id of 65536 + e is not local, and the block kernel says so; the + cluster kernel reads ids through int16 and sorts it as expert e.""" + for tokens in (2, 300): + ids = _ids("router", tokens, 128, seed=33) + ids[0, 0] = 40000 # negative as int16: non-local everywhere + out, _ = _check(_args(ids, 128)) + assert int(out[2][0, 0]) == -1 + ids[0, 0] = 65536 + 3 + out = op(**_args(ids, 128)) + p = int(out[2][0, 0]) + if tokens <= 4: + assert p == -1 + _compare(out, op.reference(**_args(ids, 128))) + else: + assert p >= 0 and int(out[0][p // 128]) == 3 + with pytest.raises(AssertionError): + _compare(out, op.reference(**_args(ids, 128))) + + +# ── loud rejections ─────────────────────────────────────────────────────── + + +def test_op_rejects_loudly(): + """Rejected by the op itself; none of these needs a guard.""" + ids = _ids("router", 64, 128, seed=41) + sc = torch.rand(64, K, device=DEV) + wide33 = _ids("router", 8, 128, seed=42, top_k=33) + cases = [ + ("int64 ids", lambda: op(ids.long(), sc, E, K, 0, E, 128), "expected scalar type Int"), + ("1-D ids", lambda: op(ids.reshape(-1), sc, E, K, 0, E, 128), "must be 2D"), + ("top_k vs width", lambda: op(ids, sc, E, 9, 0, E, 128), r"size\(1\) must be top_k"), + ("scales width", lambda: op(ids, sc[:, :9], E, K, 0, E, 128), r"size\(1\) must be top_k"), + ("scales rows", lambda: op(ids, sc[:63], E, K, 0, E, 128), "must be num_tokens"), + ("top_k 33", lambda: op(wide33, torch.rand(8, 33, device=DEV), E, 33, 0, E, 128), "<= 32"), + ("num_experts 2049", lambda: op(ids, sc, 2049, K, 0, 2049, 128), "no more than 2048"), + ("T = 0", lambda: op(ids[:0], sc[:0], E, K, 0, E, 128), "mPtrTopKWeights must also"), + ] + for tag, call, match in cases: + with pytest.raises(RuntimeError, match=match): + call() + with pytest.raises(NotImplementedError): + op(ids.cpu(), sc.cpu(), E, K, 0, E, 128) + + +# ── the guards ──────────────────────────────────────────────────────────── + + +def test_unguarded_noncontiguous_ids_answer_wrongly(): + """What `is_valid` exists for: the op reads a strided view as the dense buffer + that starts at its data pointer, and answers without raising.""" + wide = _ids("router", 300, 128, seed=51, top_k=16) + view = wide[:, :K] + sc = torch.rand(300, K, device=DEV) + out = op(view, sc, E, K, 0, E, 128) + dense = wide.reshape(-1)[: 300 * K].reshape(300, K) + _compare(out, op.reference(dense, sc, E, K, 0, E, 128)) + with pytest.raises(AssertionError): + _compare(out, op.reference(view.contiguous(), sc, E, K, 0, E, 128)) + + +def test_guards_refuse(): + """The 2^31-element operands are expanded views of one row: the guard reads + only their shape, so nothing is allocated.""" + wide = _ids("router", 300, 128, seed=51, top_k=16) + sc = torch.rand(300, K, device=DEV) + ids = _ids("router", 300, 128, seed=53) + base = ids[:1] + with validating(op): + with pytest.raises(AssertionError, match="contiguous"): + op(wide[:, :K], sc, E, K, 0, E, 128) + with pytest.raises(AssertionError, match="contiguous"): + op(ids.t().contiguous().t(), sc, E, K, 0, E, 128) + for tokens in (214748365, 429496730): + with pytest.raises(AssertionError, match=r"2\^31"): + op(base.expand(tokens, K), sc[:1].expand(tokens, K), E, K, 0, E, 128) + # tile_tokens_dim = 0 is refused but not driven unguarded: the host divides + # by it, which returns garbage on aarch64 and traps on x86. + with pytest.raises(AssertionError, match="tile_tokens_dim must be positive"): + op(ids, sc, E, K, 0, E, 0) + out = op(wide[:, :K].contiguous(), sc, E, K, 0, E, 128) + _compare(out, op.reference(wide[:, :K].contiguous(), sc, E, K, 0, E, 128)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_unpermute_inplace.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_unpermute_inplace.py new file mode 100644 index 000000000000..d46140fabf49 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_unpermute_inplace.py @@ -0,0 +1,527 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the moe_unpermute_inplace entry of the Qwen3.8-Flash-Next target. + +The entry owns the contract. This file lays a cell's FC2 rows out the way the +CuteDSL FC2 hands them over -- placed by moe_sort, every row the op must not read +filled with NaN, the output pre-filled with NaN -- and holds the claims no cell +carries: CUDA-graph replay at the captured row counts, the eager prefix sweep, +run-to-run and placement independence, the gate's discrimination, T = 0, the loud +rejections 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.moe_unpermute_inplace import ( # noqa: E501 + moe_unpermute_inplace as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "moe_unpermute_inplace requires a CUDA device" + +DEV = "cuda" +E, K, H = 512, 10, 2560 +NAN = float("nan") + +#: Row counts of a decode capture (requests x MTP verify width 1..4, up to 32 +#: requests), and the eager prefix lengths: both sides of the power-of-two edges up +#: to the 8192-row ceiling. A subset of the prefixes is also re-sorted. +_CAPTURED_ROWS = (1, 2, 3, 4, 5, 8, 32, 96, 128) +_PREFIXES = ( + *(1, 2, 3, 4, 5, 8, 31, 32, 33, 96, 127, 128, 129, 511, 512, 513, 1023, 1024, 1025), + *(2047, 2048, 2049, 4095, 4096, 4097, 8191), +) +_RESORT = (1, 2, 5, 32, 129, 1024, 8192) + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Bitwise. + + `reference` is the computation the op performs, in its order and with its + roundings, so no band is needed and none is given. The gate sees the + accumulation order: the same sum taken in reverse slot order, with each + product rounded before its add, as the float64 sum rounded once, or + accumulated in bf16, differs in many elements and is rejected. A failure in a + few isolated one-ulp elements points at the kernel's rounding (a lost + multiply-add contraction, say); a wrong row or scale moves whole rows. + """ + assert out.dtype == ref.dtype, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (tuple(out.shape), tuple(ref.shape)) + bits = {2: torch.int16, 4: torch.int32}[out.element_size()] + differ = out.view(bits) != ref.view(bits) + n = int(differ.sum()) + if n: + first = torch.nonzero(differ)[0].tolist() + raise AssertionError( + f"{n} of {out.numel()} elements differ bitwise from the slot-order fp32 " + f"sum; first at {first}: {out[tuple(first)].item()} vs " + f"{ref[tuple(first)].item()}" + ) + + +def _spec(tokens, tile_tokens_dim=128, dist="router", values="normal", non_local_from=None): + """The fields `_problem` reads; the same names as a cell's spec.""" + return dict( + tokens=tokens, + tile_tokens_dim=tile_tokens_dim, + dist=dist, + values=values, + non_local_from=non_local_from, + ) + + +def _ids(dist: str, tokens: int, seed: int) -> torch.Tensor: + """Expert ids, `[tokens, K]` int32, every row naming 10 distinct experts.""" + g = torch.Generator(device=DEV).manual_seed(seed) + if dist == "router": + scores = torch.rand(tokens, E, device=DEV, generator=g) + return scores.topk(K, dim=-1).indices.to(torch.int32) + if dist == "hot": + hot = torch.randperm(E, device=DEV, generator=g)[:K] + slots = torch.rand(tokens, K, device=DEV, generator=g).argsort(dim=-1) + return hot[slots].to(torch.int32) + if dist == "round_robin": + base = torch.arange(tokens, device=DEV).unsqueeze(1) * K + return ((base + torch.arange(K, device=DEV)) % E).to(torch.int32) + raise ValueError(dist) + + +def _scales(tokens: int, seed: int) -> torch.Tensor: + """fp32 routing weights, a softmax over each row's ten slots.""" + g = torch.Generator(device=DEV).manual_seed(seed + 1) + return torch.softmax(torch.randn(tokens, K, device=DEV, generator=g), dim=-1) + + +def _rows(tokens: int, seed: int, values: str = "normal") -> torch.Tensor: + """The FC2 output row of each (token, slot) pair, `[tokens, K, H]` bf16.""" + g = torch.Generator(device=DEV).manual_seed(seed + 2) + v = torch.randn(tokens, K, H, device=DEV, generator=g) + if values == "wide": + v = v * torch.exp2(torch.randint(-30, 31, v.shape, device=DEV, generator=g).float()) + elif values != "normal": + raise ValueError(values) + return v.to(torch.bfloat16) + + +def _sort(ids: torch.Tensor, scales: torch.Tensor, tile: int): + """moe_sort as the CuteDSL MoE calls it: (expanded_idx_to_permuted_idx, padded rows).""" + out = torch.ops.trtllm.moe_sort( + token_selected_experts=ids, + token_final_scales=scales, + num_experts=E, + top_k=K, + local_expert_offset=0, + local_num_experts=E, + tile_tokens_dim=tile, + ) + return out[2], out[3].numel() + + +def _place(rows: torch.Tensor, e2p: torch.Tensor, num_rows: int) -> torch.Tensor: + """The FC2 output as the op receives it: each live pair's row at its permuted + index, NaN in every row the op has no business reading.""" + perm = torch.full((num_rows, H), NAN, dtype=rows.dtype, device=DEV) + flat = e2p.reshape(-1) + live = flat >= 0 + perm[flat[live].long()] = rows.reshape(-1, H)[live] + return perm + + +def _output(tokens: int, dtype: torch.dtype = torch.bfloat16) -> torch.Tensor: + return torch.full((tokens, H), NAN, dtype=dtype, device=DEV) + + +def _problem(spec: dict, seed: int): + """A cell's arguments, plus the per-pair rows and ids they were built from.""" + tokens = spec["tokens"] + ids = _ids(spec["dist"], tokens, seed) + scales = _scales(tokens, seed) + e2p, num_rows = _sort(ids, scales, spec["tile_tokens_dim"]) + if spec["non_local_from"] is not None: + e2p = torch.where(ids >= spec["non_local_from"], -1, e2p).to(torch.int32) + e2p[0] = -1 + rows = _rows(tokens, seed, spec["values"]) + args = dict( + permuted_input=_place(rows, e2p, num_rows), + output=_output(tokens), + expanded_idx_to_permuted_idx=e2p, + topk_scales=scales, + ) + return args, rows, ids + + +def _bits(t: torch.Tensor) -> torch.Tensor: + return t.contiguous().view(torch.uint8) + + +def _same_bits(a: torch.Tensor, b: torch.Tensor) -> bool: + return a.shape == b.shape and a.dtype == b.dtype and torch.equal(_bits(a), _bits(b)) + + +def _cell_id(cell): + s = cell.spec + tag = f"t{s['tile_tokens_dim']}-T{s['tokens']}-{s['dist']}-{s['values']}" + return tag if s["non_local_from"] is None else f"{tag}-nonlocal{s['non_local_from']}" + + +# ── the certified cells ─────────────────────────────────────────────────── + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[_cell_id(c) for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + s = cell.spec + assert (s["num_experts"], s["top_k"], s["hidden_size"]) == (E, K, H) + assert (s["dtype"], s["scales_dtype"]) == (torch.bfloat16, torch.float32) + args, _, _ = _problem(s, zlib.crc32(cell.why.encode())) + before = { + n: args[n].clone() + for n in ("permuted_input", "expanded_idx_to_permuted_idx", "topk_scales") + } + out = op(**args) + assert out is args["output"], "the op must return the tensor it wrote" + _compare(out, op.reference(**args)) + for name, snapshot in before.items(): + assert _same_bits(args[name], snapshot), f"{name} was written" + if s["non_local_from"] is not None: + assert int((args["expanded_idx_to_permuted_idx"] < 0).sum()) > s["tokens"] + assert bool((_bits(out[0]) == 0).all()), "a token with every slot -1 is not +0" + + +# ── CUDA graphs: captured row counts, on new data ───────────────────────── + + +def _capture(fn): + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for _ in range(2): + fn() + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + static = fn() + return graph, static + + +@pytest.mark.parametrize("tile", (128, 256)) +def test_captured_rows_replay_on_new_data(tile, certifying): # noqa: F811 + """The sort, the FC2 rows' placement and this op are captured together, as the + target runs them; each graph is replayed on two fresh sets of ids, scales and + rows. The permuted buffer is not cleared between replays, so the rows the new + sort does not use hold the previous placement's rows. Each replay is gated + against the reference, and is bitwise equal to an eager call on the same data + through its own sort. + """ + certifying(op) + stale_rows = 0 + for tokens in _CAPTURED_ROWS: + ids = _ids("router", tokens, seed=tokens) + scales = _scales(tokens, seed=tokens) + rows = _rows(tokens, seed=tokens) + num_rows = _sort(ids, scales, tile)[1] + perm = torch.full((num_rows, H), NAN, dtype=torch.bfloat16, device=DEV) + out = _output(tokens) + + def step(ids=ids, scales=scales, rows=rows, perm=perm, out=out, tile=tile): + e2p = _sort(ids, scales, tile)[0] + perm.index_copy_(0, e2p.reshape(-1).long(), rows.reshape(-1, H)) + op( + permuted_input=perm, + output=out, + expanded_idx_to_permuted_idx=e2p, + topk_scales=scales, + ) + return e2p + + graph, e2p = _capture(step) + for trial in (1, 2): + seed = 1000 * trial + tokens + ids.copy_(_ids("router", tokens, seed)) + scales.copy_(_scales(tokens, seed)) + rows.copy_(_rows(tokens, seed)) + out.fill_(NAN) + graph.replay() + args = dict( + permuted_input=perm, + output=out, + expanded_idx_to_permuted_idx=e2p, + topk_scales=scales, + ) + _compare(out, op.reference(**args)) + used = torch.zeros(num_rows, dtype=torch.bool, device=DEV) + used[e2p.reshape(-1).long()] = True + stale_rows += int((~used & ~perm.isnan().any(1)).sum()) + + e2p_eager, n = _sort(ids.clone(), scales.clone(), tile) + eager = op( + permuted_input=_place(rows.clone(), e2p_eager, n), + output=_output(tokens), + expanded_idx_to_permuted_idx=e2p_eager, + topk_scales=scales.clone(), + ) + assert _same_bits(out, eager), (tokens, trial) + del graph, e2p + torch.cuda.synchronize() + assert stale_rows > 0, "no replay ran over a stale row: the stale-row claim is untested" + + +# ── eager: prefixes of the largest batch ────────────────────────────────── + + +@pytest.mark.parametrize("tile", (128, 256)) +def test_eager_rows_prefix_bitwise(tile, certifying): # noqa: F811 + """One block per token makes the grid the only thing T sizes: the 8192-row call + is gated against the reference, and the first T rows of every row count in the + sweep are bitwise those of the 8192-row call -- over the same placement, and + (a subset) over a fresh sort of the first T tokens alone, which places their + rows elsewhere. A token's output does not depend on its batch. + """ + certifying(op) + args, rows, ids = _problem(_spec(8192, tile_tokens_dim=tile), seed=8192 + tile) + full = op(**args) + _compare(full, op.reference(**args)) + perm, e2p, scales = ( + args["permuted_input"], + args["expanded_idx_to_permuted_idx"], + args["topk_scales"], + ) + for tokens in _PREFIXES: + out = op( + permuted_input=perm, + output=_output(tokens), + expanded_idx_to_permuted_idx=e2p[:tokens], + topk_scales=scales[:tokens], + ) + assert _same_bits(out, full[:tokens]), tokens + for tokens in _RESORT: + e2p_t, n = _sort(ids[:tokens], scales[:tokens], tile) + out = op( + permuted_input=_place(rows[:tokens], e2p_t, n), + output=_output(tokens), + expanded_idx_to_permuted_idx=e2p_t, + topk_scales=scales[:tokens], + ) + assert _same_bits(out, full[:tokens]), tokens + + +# ── run to run, and where the rows sit ──────────────────────────────────── + + +@pytest.mark.parametrize("tokens", (1, 32, 128, 8192)) +def test_repeats_and_placements_are_bitwise(tokens, certifying): # noqa: F811 + """What the fused finalize does not give: 8 identical calls, a second sort + of the same ids (the sort orders an expert's rows by atomics), a random + permutation of every permuted row, and a shuffle of the tokens all give the + same bits. + """ + certifying(op) + args, rows, ids = _problem(_spec(tokens), seed=tokens + 7) + first = op(**args).clone() + _compare(first, op.reference(**args)) + for _ in range(7): + assert _same_bits(op(**dict(args, output=_output(tokens))), first) + + perm, e2p, scales = ( + args["permuted_input"], + args["expanded_idx_to_permuted_idx"], + args["topk_scales"], + ) + e2p2, n = _sort(ids, scales, 128) + if tokens == 8192: + assert int((e2p2 != e2p).sum()) > 0, "the second sort placed every row where the first did" + out = op( + permuted_input=_place(rows, e2p2, n), + output=_output(tokens), + expanded_idx_to_permuted_idx=e2p2, + topk_scales=scales, + ) + assert _same_bits(out, first), "a second sort of the same ids changed the result" + + g = torch.Generator(device=DEV).manual_seed(tokens) + p = torch.randperm(perm.shape[0], device=DEV, generator=g) + moved = torch.empty_like(perm) + moved[p] = perm + out = op( + permuted_input=moved, + output=_output(tokens), + expanded_idx_to_permuted_idx=p[e2p.long()].to(torch.int32), + topk_scales=scales, + ) + assert _same_bits(out, first), "a random placement of the rows changed the result" + + order = torch.randperm(tokens, device=DEV, generator=g) + out = op( + permuted_input=perm, + output=_output(tokens), + expanded_idx_to_permuted_idx=e2p[order].contiguous(), + topk_scales=scales[order].contiguous(), + ) + assert _same_bits(out, first[order]), "shuffling the tokens changed a token's result" + + +# ── what the bitwise gate can see ───────────────────────────────────────── + + +def test_the_gate_sees_slot_order_and_rounding(certifying): # noqa: F811 + """`_compare` is bitwise; a sum taken in another order or with other roundings + is rejected. The reverse-order sum is the entry's own reference over the + slots reversed. + """ + certifying(op) + args, rows, _ = _problem(_spec(8192), seed=99) + out = op(**args) + _compare(out, op.reference(**args)) + scales = args["topk_scales"] + reverse = op.reference( + permuted_input=args["permuted_input"], + output=args["output"], + expanded_idx_to_permuted_idx=args["expanded_idx_to_permuted_idx"].flip(1).contiguous(), + topk_scales=scales.flip(1).contiguous(), + ) + unfused = torch.zeros(8192, H, device=DEV) + for k in range(K): + unfused = unfused + rows[:, k].float() * scales[:, k : k + 1] + exact = (rows.double() * scales.double().unsqueeze(-1)).sum(1) + bf16_sum = torch.zeros(8192, H, dtype=torch.bfloat16, device=DEV) + for k in range(K): + bf16_sum = bf16_sum + (rows[:, k].float() * scales[:, k : k + 1]).to(torch.bfloat16) + for wrong in ( + reverse, # reverse slot order + unfused.to(torch.bfloat16), # product rounded before the add + exact.to(torch.bfloat16), # float64 sum rounded once + bf16_sum, # bf16 accumulation + ): + with pytest.raises(AssertionError, match="differ bitwise"): + _compare(out, wrong) + + +# ── T = 0 and the loud rejections ───────────────────────────────────────── + + +def test_zero_tokens_returns_and_leaves_no_error(): + args, _, _ = _problem(_spec(8), seed=71) + empty = op( + permuted_input=args["permuted_input"], + output=torch.empty(0, H, dtype=torch.bfloat16, device=DEV), + expanded_idx_to_permuted_idx=args["expanded_idx_to_permuted_idx"][:0], + topk_scales=args["topk_scales"][:0], + ) + torch.cuda.synchronize() + assert tuple(empty.shape) == (0, H) + assert float((torch.ones(1000, device=DEV) + 1).sum()) == 2000.0 + out = op(**args) + _compare(out, op.reference(**args)) + + +def test_op_rejects_loudly(): + """Rejected by the op itself; none of these needs a guard.""" + tokens = 64 + args, _, _ = _problem(_spec(tokens), seed=81) + p, e, s = ( + args["permuted_input"], + args["expanded_idx_to_permuted_idx"], + args["topk_scales"], + ) + narrow = 2556 + cases = [ + ("int64 indices", dict(expanded_idx_to_permuted_idx=e.long()), "expected scalar type Int"), + ( + "hidden not a multiple of 8", + dict( + permuted_input=p[:, :narrow].contiguous(), + output=torch.empty(tokens, narrow, dtype=torch.bfloat16, device=DEV), + ), + "hidden_size must be divisible by 8", + ), + ( + "fp32 rows", + dict(permuted_input=p.float(), output=_output(tokens, torch.float32)), + "Unsupported input dtype", + ), + ("fp16 scales, bf16 rows", dict(topk_scales=s.half()), "Unsupported input dtype"), + ( + "fewer than T * top_k rows", + dict(permuted_input=p[: tokens * K - 1]), + r"greater than or equal to num_tokens \* top_k", + ), + ("scales rows", dict(topk_scales=s[:-1]), r"topk_scales.size\(0\) must be num_tokens"), + ( + "scales width", + dict(topk_scales=s[:, :9].contiguous()), + r"topk_scales.size\(1\) must be top_k", + ), + ( + "index rows", + dict(expanded_idx_to_permuted_idx=e[:-1]), + r"expanded_idx_to_permuted_idx.size\(0\) must be num_tokens", + ), + ( + "output width", + dict(output=torch.empty(tokens, H - 8, dtype=torch.bfloat16, device=DEV)), + r"output.size\(1\) must be hidden_size", + ), + ] + for tag, change, match in cases: + with pytest.raises(RuntimeError, match=match): + op(**dict(args, **change)) + torch.cuda.synchronize() + out = op(**args) + _compare(out, op.reference(**args)) + + +# ── the guards ──────────────────────────────────────────────────────────── + +_OPERANDS = ("permuted_input", "output", "expanded_idx_to_permuted_idx", "topk_scales") + + +def _strided(t: torch.Tensor, pad: int) -> torch.Tensor: + """The same values in a view whose rows are `pad` elements wider apart.""" + wide = torch.zeros(t.shape[0], t.shape[1] + pad, dtype=t.dtype, device=t.device) + wide[:, : t.shape[1]] = t + return wide[:, : t.shape[1]] + + +@pytest.mark.parametrize("name", _OPERANDS) +def test_noncontiguous_operand_answers_wrongly_and_is_refused(name): + args, _, _ = _problem(_spec(300), seed=61) + ref = op.reference(**args) + bad = dict(args, **{name: _strided(args[name], 64 if args[name].shape[1] == H else 6)}) + out = op(**bad) + with pytest.raises(AssertionError, match="differ bitwise"): + _compare(out.contiguous(), ref) + with validating(op): + with pytest.raises(AssertionError, match=f"{name} must be contiguous"): + op(**bad) + out = op(**dict(args, output=_output(300))) + _compare(out, ref) + + +def test_output_dtype_mismatch_answers_wrongly_and_is_refused(): + """An fp32 output receives the bf16 result's bit patterns, packed into the + first half of its bytes. + """ + args, _, _ = _problem(_spec(300), seed=62) + ref = op.reference(**args) + out32 = op(**dict(args, output=_output(300, torch.float32))) + assert _same_bits(out32.view(torch.bfloat16).reshape(-1)[: 300 * H].view(300, H), ref) + with pytest.raises(AssertionError): + _compare(out32, ref) + with validating(op): + with pytest.raises(AssertionError, match="output dtype"): + op(**dict(args, output=_output(300, torch.float32))) + + +def test_tokens_past_int32_are_refused(): + """The guard reads only `output.size(0)` first, so a meta tensor stands in for + the 2^31-row output the op would silently leave unwritten.""" + args, _, _ = _problem(_spec(8), seed=63) + huge = torch.empty((2**31, 8), dtype=torch.bfloat16, device="meta") + with validating(op), pytest.raises(AssertionError, match=r">= 2\^31 rows"): + op(**dict(args, output=huge)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_triton_update.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_triton_update.py index 2d465304f39b..c9ef151454bf 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_triton_update.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_triton_update.py @@ -80,7 +80,7 @@ def randn(*shape): rows = spec.get("snapshots", 0) window = rows_idx = None if rows: - steps = 1 if x_form == "2d" else seqlen + steps = spec.get("snapshot_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 [ diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_chunk_gated_delta_rule.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_chunk_gated_delta_rule.py index 92fff412c4f4..678033a7a15a 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_chunk_gated_delta_rule.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_chunk_gated_delta_rule.py @@ -4,8 +4,9 @@ The first indexed in-place call per autotune key in a process must be as correct as a warm one: the state kernel's autotuner restores the indexed state its benchmark launches -write. "First call in the process" is a property of the process, so the cold-call tests -re-run this file as a child: `python `. +write, and the pool's dtype is part of the key. "First call in the process" is a property +of the process, so the cold-call tests re-run this file as a child: +`python `. """ import os @@ -110,11 +111,13 @@ def _operands( return q, k, v, g, beta, cu_seqlens -def _pool(num_slots: int, seed: int, num_v_heads: int = NUM_V_HEADS) -> torch.Tensor: - """An fp32 recurrent-state pool, `[slots, H, V, K]` (K innermost).""" +def _pool( + num_slots: int, seed: int, num_v_heads: int = NUM_V_HEADS, dtype: torch.dtype = torch.float32 +) -> torch.Tensor: + """A recurrent-state pool, `[slots, H, V, K]` (K innermost), at `dtype`.""" gen = torch.Generator(device=DEV).manual_seed(seed) shape = (num_slots, num_v_heads, HEAD_V_DIM, HEAD_K_DIM) - return torch.randn(shape, generator=gen, device=DEV) * 0.1 + return (torch.randn(shape, generator=gen, device=DEV) * 0.1).to(dtype) def _gate_output(o: torch.Tensor, ref: tuple) -> None: @@ -137,15 +140,16 @@ def _build(spec, seed): normalize_qk=spec.get("normalize_qk", True), ) num_seqs = q.shape[0] if cu_seqlens is None else len(cu_seqlens) - 1 + pool_dtype = spec.get("pool_dtype", torch.float32) initial, indices, inplace = None, None, False if spec["initial"] == "indexed": - initial = _pool(spec["slots"], seed=seed + 1) + initial = _pool(spec["slots"], seed=seed + 1, dtype=pool_dtype) 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) + initial = _pool(num_seqs, seed=seed + 1, dtype=pool_dtype) wide, output = None, None shape_o = (q.shape[0], q.shape[1], NUM_V_HEADS, HEAD_V_DIM) if spec.get("output") == "buffer": @@ -306,9 +310,9 @@ def _state_kernel_autotuner(): raise AssertionError("no triton autotuner found around the chunked state kernel") -def _checkpoint_case(seed: int = 101): +def _checkpoint_case(seed: int = 101, pool_dtype: torch.dtype = torch.float32): q, k, v, g, beta, cu = _operands(CHILD_SEQLENS, seed=seed) - pool = _pool(4, seed=seed + 1) + pool = _pool(4, seed=seed + 1, dtype=pool_dtype) indices = torch.tensor([2, 0], dtype=torch.int32, device=DEV) buf = torch.empty( 1, sum(CHILD_SEQLENS), NUM_V_HEADS, HEAD_V_DIM, device=DEV, dtype=torch.bfloat16 @@ -316,10 +320,10 @@ def _checkpoint_case(seed: int = 101): return (q, k, v, g, beta, None, pool, indices, True, False, cu, False, buf) -def _child_cold_call() -> None: +def _child_cold_call(pool_dtype: torch.dtype = torch.float32) -> None: autotuner = _state_kernel_autotuner() assert not autotuner.cache, "the child must be the first indexed call in its process" - args = _checkpoint_case() + args = _checkpoint_case(pool_dtype=pool_dtype) ref = op.reference(*args) o, _ = op(*args) assert len(autotuner.cache) == 1, "the call must have benchmarked exactly one key" @@ -344,6 +348,7 @@ def _child_cold_call_odd_heads() -> None: _CHILD_MODES = { "cold-call": _child_cold_call, + "cold-call-bf16-pool": lambda: _child_cold_call(torch.bfloat16), "cold-call-odd-heads": _child_cold_call_odd_heads, } @@ -374,6 +379,11 @@ def test_cold_indexed_call_at_this_checkpoints_geometry() -> None: _run_child("cold-call") +def test_cold_indexed_call_over_a_bf16_pool() -> None: + """The same call over a bf16 pool, the autotune key an fp32 pool never touches.""" + _run_child("cold-call-bf16-pool") + + 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") diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_cu_seqlens_from_seq_lens.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_cu_seqlens_from_seq_lens.py index 5e011da0ac63..5fcd1d61448b 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_cu_seqlens_from_seq_lens.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_cu_seqlens_from_seq_lens.py @@ -132,10 +132,12 @@ def test_accumulation_happens_in_the_requested_dtype() -> None: assert op(seq_lens, torch.int32).tolist() == [0, 2_000_000_000, -294_967_296] -def test_cuda_graph_capture_and_replay() -> None: +@pytest.mark.parametrize("batch, p", ((8, 2), (32, 8))) +def test_cuda_graph_capture_and_replay(batch, p) -> None: """A captured graph follows lengths rewritten in place: only the shape - reaches Python, the values are read on device and never synchronized on.""" - static_lens = _seq_lens([2] * 8) + reaches Python, the values are read on device and never synchronized on. Eight + requests of 2 rows, and 32 of 8 (the MTP-7 verify's 256 rows).""" + static_lens = _seq_lens([p] * batch) side = torch.cuda.Stream() side.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(side): @@ -147,7 +149,7 @@ def test_cuda_graph_capture_and_replay() -> None: graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): static_out = op(static_lens, torch.int32) - for lengths in ([1] * 8, [2] * 8, [3, 1, 2, 1, 1, 1, 1, 1]): + for lengths in ([1] * batch, [p] * batch, [p + 1] + [1] * (batch - 1)): static_lens.copy_(_seq_lens(lengths)) graph.replay() torch.cuda.synchronize() diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_dilated_short_conv_paged.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_dilated_short_conv_paged.py index a7d49b0d7aa5..7233481e3ef2 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_dilated_short_conv_paged.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_dilated_short_conv_paged.py @@ -333,9 +333,11 @@ def test_repeat_calls_are_bitwise_stable() -> None: 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) +@pytest.mark.parametrize("batch, p", ((4, 4), (32, 8))) +def test_cuda_graph_capture_and_replay(batch, p) -> None: + """Captured once, replayed after the batch tensors' values change in place: four + requests of 4 tokens, and the MTP-7 verify's 32 of 8.""" + args = _small((p,) * batch, tuple(range(batch)), seed=25, pool_slots=batch + 2) x, weight, _, pool, qsl, cidx, his = args[:7] live = pool.clone() stream = torch.cuda.Stream() @@ -351,9 +353,10 @@ def test_cuda_graph_capture_and_replay() -> None: captured = op(x, weight, None, live, qsl, cidx, his, DILATION, True, PAD_SLOT_ID) g = torch.Generator(device="cuda").manual_seed(26) - qsl.copy_(torch.tensor([0, 1, 2, 3, 16], dtype=torch.int32, device="cuda")) - cidx.copy_(torch.tensor([4, 3, 2, 1], dtype=torch.int32, device="cuda")) - his.copy_(torch.tensor([True, False, True, False], dtype=torch.bool, device="cuda")) + # Every row but the last shrinks to one token; the slot map reverses and shifts by one. + qsl.copy_(torch.tensor([*range(batch), batch * p], dtype=torch.int32, device="cuda")) + cidx.copy_(torch.tensor(list(range(batch, 0, -1)), dtype=torch.int32, device="cuda")) + his.copy_(torch.tensor([i % 2 == 0 for i in range(batch)], dtype=torch.bool, device="cuda")) x.copy_(torch.randn(x.shape, generator=g, device="cuda").to(x.dtype)) live.copy_(pool) graph.replay() diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py new file mode 100644 index 000000000000..b47fd95fa910 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py @@ -0,0 +1,884 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the fused_recurrent_gated_delta_rule_cached_replay_update entry. + +The op is stateful across verify steps: what one call leaves in the checkpoint +pool and the history double buffer is what the next call reads, after the +engine's own bookkeeping ran in between. So every cell drives several +consecutive verify steps the way the engine does -- work items before the +forward (partitioned batches), the op once per layer, then the all-layer +commit and `_advance_replay_state` with a random partial acceptance -- and +checks two things at every step: + +* each call against the entry's `reference` on that call's own inputs (the + output and all four state buffers), plus bitwise locks on everything the + call must not touch; +* every output against an *exact twin*: the entry's `reference` driven on the + same checkpoint and history held in fp32 and deep enough never to wrap, + advanced by `p += accepted` alone. The twin's state is the plain + token-by-token recurrence over every accepted token, independent of the + wrap, the double buffer and the commit, so a wrong commit or wrong + bookkeeping surfaces as an output that drifts off it -- + `test_the_twin_sees_a_broken_protocol` shows it does. +""" + +from __future__ import annotations + +import zlib + +import pytest +import torch +import torch.nn.functional as F +from triton.compiler.errors import CompilationError + +import tensorrt_llm._torch.modules.fla.cached_replay 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_cached_replay_update import ( # noqa: E501 + ALL_LAYER_COMMIT_MIN_BATCH, + HEAD_K_DIM, + HEAD_V_DIM, + HISTORY_SIZE, + NUM_K_HEADS, + NUM_V_HEADS, + PACKED_QKV_WIDTH, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_recurrent_gated_delta_rule_cached_replay_update import ( # noqa: E501 + fused_recurrent_gated_delta_rule_cached_replay_update as op, +) +from tensorrt_llm._torch.modules.fla.cached_replay import ( + CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE, + commit_gdn_cached_replay_history_layers, +) +from tensorrt_llm._torch.modules.mamba.mamba2_metadata import _build_replay_work_items_triton +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( + ReplayStateUpdateMetadata, + _advance_replay_state, +) + +__extra_import_path__ = ["../../../.."] +from _validating import ULP, certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), ( + "fused_recurrent_gated_delta_rule_cached_replay_update requires a CUDA device" +) + +DEV = "cuda" +H, HV, K, V = NUM_K_HEADS, NUM_V_HEADS, HEAD_K_DIM, HEAD_V_DIM +W = PACKED_QKV_WIDTH +HIST = HISTORY_SIZE +OUT_SENTINEL = -7.0 +BF16_ULP = ULP[torch.bfloat16] + +#: This checkpoint's 36 GDN layers of `linear_attn.A_log` / `linear_attn.dt_bias` +#: (bf16, 48 each): mean, standard deviation, range. +A_LOG_STATS = (1.006, 1.487, (-7.125, 5.0625)) +DT_BIAS_STATS = (-4.411, 1.732, (-11.4375, 2.53125)) + +#: Element gate of the bf16 members, in bf16 ulps of the block they belong to. +_ELEMENT_ULPS = 4.0 +#: Relative-RMS gate of the output, in bf16 ulps. +_OUTPUT_RMS_ULPS = 1.5 +#: Element gate of an fp32 checkpoint state, relative to its (slot, head) scale. +_STATE_TOL = 2.0**-13 +#: Element gate of a bf16 checkpoint state, in bf16 ulps of its (slot, head) scale. +_BF16_STATE_ULPS = 2.0 +#: Gate of the stored cumulative log-decays: |dG| <= atol + rtol * |G|. +_G_ATOL, _G_RTOL = 2.0**-11, 2.0**-18 + + +def _gate(name, got, want, element, rms=None): + """`got` against `want` per block of the last two dims (a token's + `[heads, dim]`, or one head's `[V, K]` state): every element within + `element` of its block's absmax, NaN blocks of `want` excluded.""" + block = got.shape[-2] * got.shape[-1] + g, w = got.float().reshape(-1, block), want.float().reshape(-1, block) + keep = ~torch.isnan(w).any(1) + g, w = g[keep], w[keep] + scale = w.abs().amax(1, keepdim=True).clamp_min(torch.finfo(torch.float32).tiny) + ratio = ((g - w).abs() / scale).nan_to_num(nan=float("inf")) + worst = ratio.max().item() if ratio.numel() else 0.0 + assert worst <= element, ( + f"{name}: an element is {worst:.3e} of its block's absmax from the reference " + f"(gate {element:.3e})" + ) + if rms is not None and w.numel(): + rel = ((g - w).pow(2).mean().sqrt() / w.pow(2).mean().sqrt().clamp_min(1e-30)).item() + assert rel <= rms, f"{name}: relative RMS {rel:.3e} exceeds {rms:.3e}" + + +def _compare(out: tuple, ref: tuple) -> None: + """`(returned, ssm_states, old_u, old_k, old_G)` as the call left them, + against `reference`'s five; a member `None` on both sides is skipped (the + exact twin is compared on the output alone). + + The kernel is not fp32 throughout, and the bands say where. It rounds the + L2-normalized q and k to bf16 before every product (2^-9 relative each), + rounds the history-mixing coefficients to bf16 for its tensor-core dots, and + stores the output and the history's u and k in bf16. So the output, `old_u` + and `old_k` are gated per block -- one token's `[heads, dim]` -- at 4 bf16 + ulps of the block's absmax, which keeps a near-silent head from turning + rounding into relative noise, and the output's relative RMS at 1.5 ulps. An fp32 + checkpoint state is fp32 by construction: the commit splits the decayed + history update into a bf16 hi and lo part, ~16 significant bits per term, and + up to 16 terms each within 2^-17 give the 2^-13 gate, relative to one (slot, + head) state's absmax. A bf16 checkpoint is that fp32 commit rounded once, as the + reference's is, so the two stored values are equal or adjacent, and adjacent bf16 + values are at most 2 bf16 ulps of their head's absmax apart: the bf16 gate is one + representable step. `G` only ever enters as exp(G_t - G_j), so an absolute + error is a relative decay error: the kernel's softplus runs on approximate + transcendentals whose fp32 sum drops exp(x)'s low bits; atol 2^-11 is an + eighth of a bf16 ulp of the u it scales, and rtol 2^-18 covers the cumulative + sum's rounding where |G| reaches thousands on fast-forgetting heads. + """ + names = ("output", "ssm_states", "old_u", "old_k", "old_G") + assert len(out) == len(ref) == len(names) + for name, got, want in zip(names, out, ref): + assert (got is None) == (want is None), f"{name}: present on one side only" + if got is None: + continue + assert got.dtype == want.dtype and got.shape == want.shape, ( + name, + (got.dtype, tuple(got.shape)), + (want.dtype, tuple(want.shape)), + ) + if name == "old_G": + torch.testing.assert_close(got, want, rtol=_G_RTOL, atol=_G_ATOL, msg=name) + elif name == "ssm_states": + bf16 = got.dtype == torch.bfloat16 + _gate(name, got, want, _BF16_STATE_ULPS * BF16_ULP if bf16 else _STATE_TOL) + else: + rms = _OUTPUT_RMS_ULPS * BF16_ULP if name == "output" else None + _gate(name, got, want, _ELEMENT_ULPS * BF16_ULP, rms) + + +def _seed(why): + return zlib.crc32(why.encode()) + + +def _gating(gen, layers): + """Per-layer A_log / dt_bias drawn around the checkpoint's own statistics; + head 0 pinned to the corner that forgets everything every token + (g ~ -400), head 1 to the one that hardly decays (g ~ -1e-8).""" + (a_mean, a_std, a_range), (d_mean, d_std, d_range) = A_LOG_STATS, DT_BIAS_STATS + a_log = torch.randn(layers, HV, generator=gen, device=DEV) * a_std + a_mean + dt_bias = torch.randn(layers, HV, generator=gen, device=DEV) * d_std + d_mean + a_log, dt_bias = a_log.clamp(*a_range), dt_bias.clamp(*d_range) + a_log[:, 0], dt_bias[:, 0] = a_range[1], d_range[1] + a_log[:, 1], dt_bias[:, 1] = a_range[0], d_range[0] + return a_log.contiguous(), dt_bias.contiguous() + + +def _unit(x): + return (x / x.norm(dim=-1, keepdim=True)).to(torch.bfloat16) + + +class _Bench: + """The state the cache manager keeps for `layers` GDN layers, laid out as + `MambaHybridCacheManagerV2` lays it out, driven the way the engine drives + it: work items before the forward (`_build_replay_work_items_triton`, what + mamba metadata's prepare runs for a partitioned batch), the op per layer, + then `update_mamba_states`'s order -- the all-layer commit, then + `_advance_replay_state`. + + With `exact=True` it also carries the exact twin (see the module docstring). + """ + + def __init__(self, spec, seed, exact=True): + self.spec = spec + self.path, self.n, self.t = spec["path"], spec["n"], spec["t"] + self.layers = spec.get("layers", 2) + self.slots = spec.get("slots", 64) + self.steps = spec.get("steps", 8) + self.gen = torch.Generator(device=DEV).manual_seed(seed) + self.cpu = torch.Generator().manual_seed(seed) + n, t, L, S = self.n, self.t, self.layers, self.slots + + if spec.get("edge_slots"): + inner = torch.randperm(S - 2, generator=self.cpu)[: n - 2] + 1 + chosen = torch.cat([torch.tensor([S - 1]), inner, torch.tensor([0])]) + else: + chosen = torch.randperm(S, generator=self.cpu)[:n] + chosen = chosen.to(torch.int32) + for row in spec.get("pad", ()): + chosen[row] = -1 + live_rows = chosen >= 0 + dup = tuple(spec.get("dup", ())) + if dup: + chosen[list(dup)] = int(chosen[dup[0]]) + # The padding rows' shared slot is nobody's state: the bookkeeping never names it. + live_rows[list(dup)] = False + self.state_indices = chosen.to(DEV) + self.live = self.state_indices[live_rows.to(DEV)].clone() + + # [slots, layers, HV, V, K]: each layer's state is a slot-strided view, + # the affine layout whose strides the all-layer commit is handed. + self.backing = (self._randn(S, L, HV, V, K) * 0.5).to( + spec.get("state_dtype", torch.float32) + ) + self.pool = [self.backing[:, i] for i in range(L)] + self.old_x = torch.zeros(L, S, 2, HIST, HV, V, device=DEV, dtype=torch.bfloat16) + self.old_B = torch.zeros(L, S, 2, HIST, H, K, device=DEV, dtype=torch.bfloat16) + self.old_dt = torch.zeros(L, S, 2, HV, HIST, device=DEV) + self.old_dA = self._randn(L, S, 2, HV, HIST) # handed over as old_beta; never read + self.pnat = torch.zeros(S, dtype=torch.int32, device=DEV) + self.buf = torch.zeros(S, dtype=torch.int32, device=DEV) + self.a_log, self.dt_bias = _gating(self.gen, L) + if spec.get("start", "warm") == "warm": + self.old_x.copy_(self._randn(L, S, 2, HIST, HV, V) * 0.5) + self.old_B.copy_(_unit(self._randn(L, S, 2, HIST, H, K))) + raw_a = self._randn(L, S, 2, HV, HIST) + g = -torch.exp(self.a_log)[:, None, None, :, None] * F.softplus( + raw_a + self.dt_bias[:, None, None, :, None] + ) + self.old_dt.copy_(g.cumsum(-1)) + pnat = torch.randint(0, HIST + 1, (S,), generator=self.cpu, dtype=torch.int32) + live = self.live.long().cpu() + pnat[live[0]] = HIST # a wrap at step 0 + if live.numel() > 1: + pnat[live[1]] = 0 # an append at step 0 + self.pnat.copy_(pnat) + self.buf.copy_(torch.randint(0, 2, (S,), generator=self.cpu, dtype=torch.int32)) + self.meta = ReplayStateUpdateMetadata(self.pnat, self.buf, t, HIST) + self.work_items = self.n_writes = None + if self.path == "all_layer": + self.work_items = torch.zeros(n, 4, dtype=torch.int32, device=DEV) + self.n_writes = torch.zeros(1, dtype=torch.int32, device=DEV) + + self.exact = exact + if exact: + self.depth = HIST + (self.steps + 1) * t + self.x_pool = [p.clone() for p in self.pool] + self.x_u = torch.zeros(L, S, 1, self.depth, HV, V, device=DEV) + self.x_k = torch.zeros(L, S, 1, self.depth, H, K, device=DEV) + self.x_G = torch.zeros(L, S, 1, HV, self.depth, device=DEV) + self.x_pnat = self.pnat.clone() + self.x_buf = torch.zeros_like(self.buf) + for s in self.live.tolist(): + p, b = int(self.pnat[s]), int(self.buf[s]) + self.x_u[:, s, 0, :p] = self.old_x[:, s, b, :p].float() + self.x_k[:, s, 0, :p] = self.old_B[:, s, b, :p].float() + self.x_G[:, s, 0, :, :p] = self.old_dt[:, s, b, :, :p] + + def _randn(self, *shape, dtype=torch.float32): + return torch.randn(*shape, generator=self.gen, device=DEV, dtype=dtype) + + def inputs(self): + """One layer's fresh activations: post-conv `[q|k|v]` rows and the raw + `a`/`b` projections, as the verify step hands them over.""" + n, t = self.n, self.t + packed = self._randn(n * t, W, dtype=torch.bfloat16) + rows = packed.view(n, t, W) + q = rows[..., : H * K].view(n, t, H, K) + k = rows[..., H * K : 2 * H * K].view(n, t, H, K) + v = rows[..., 2 * H * K :].view(n, t, HV, V) + if self.spec.get("packed", True) is False: + packed = None # the mixed batch's form: the views alone, copied by the op + if self.spec.get("ab", "split") == "strided": + ab = self._randn(n * t, 2 * HV, dtype=torch.bfloat16) + b, a = ab[:, :HV].view(n, t, HV), ab[:, HV:].view(n, t, HV) + else: + a = self._randn(n, t, HV, dtype=torch.bfloat16) + b = self._randn(n, t, HV, dtype=torch.bfloat16) + out_buf = output = None + lo = 0 if self.spec.get("output") == "prefix" else 5 + if self.spec.get("output") in ("prefix", "suffix"): + # A slice of a longer token-major buffer: its first rows ("prefix"), + # or the rows after 5 context tokens, as a mixed batch hands it over. + out_buf = torch.full((n * t + 5, HV, V), OUT_SENTINEL, dtype=torch.bfloat16, device=DEV) + output = out_buf[lo : lo + n * t].view(n, t, HV, V) + return dict(q=q, k=k, v=v, a=a, b=b, packed=packed, output=output, out_buf=out_buf, lo=lo) + + def kwargs(self, layer, act): + return dict( + q=act["q"], + k=act["k"], + v=act["v"], + g=act["a"], + beta=act["b"], + ssm_states=self.pool[layer], + state_indices=self.state_indices, + old_u=self.old_x[layer], + old_k=self.old_B[layer], + old_G=self.old_dt[layer], + old_beta=self.old_dA[layer], + cache_buf_idx=self.buf, + prev_num_accepted_tokens=self.pnat, + history_size=HIST, + scale=None, + use_qk_l2norm_in_kernel=True, + A_log=self.a_log[layer], + dt_bias=self.dt_bias[layer], + launch_with_pdl=self.spec.get("pdl", False), + replay_work_items=self.work_items, + n_writes=self.n_writes, + packed_qkv=act["packed"], + use_all_layer_commit=self.path == "all_layer", + output=act["output"], + ) + + def exact_kwargs(self, layer, kw): + return dict( + kw, + ssm_states=self.x_pool[layer], + old_u=self.x_u[layer], + old_k=self.x_k[layer], + old_G=self.x_G[layer], + cache_buf_idx=self.x_buf, + prev_num_accepted_tokens=self.x_pnat, + history_size=self.depth, + replay_work_items=None, + n_writes=None, + use_all_layer_commit=False, + output=None, + ) + + def adopt_exact(self, layer, ref): + assert torch.equal(ref[1], self.x_pool[layer]), "the twin must never wrap" + self.x_u[layer].copy_(ref[2]) + self.x_k[layer].copy_(ref[3]) + self.x_G[layer].copy_(ref[4]) + + def state(self, layer): + return self.pool[layer], self.old_x[layer], self.old_B[layer], self.old_dt[layer] + + def prepare(self): + if self.path == "all_layer": + self.work_items.zero_() + _build_replay_work_items_triton( + self.live, self.pnat, self.buf, self.work_items, self.n_writes, self.t, HIST + ) + + def commit(self): + if self.path == "all_layer": + commit_gdn_cached_replay_history_layers( + ssm_states=self.pool[0], + ssm_state_num_layers=self.layers, + ssm_state_layer_stride=self.backing.stride(1), + ssm_state_slot_stride=self.backing.stride(0), + old_u=self.old_x, + old_k=self.old_B, + old_G=self.old_dt, + replay_work_items=self.work_items, + n_writes=self.n_writes, + history_size=HIST, + ) + + def draw_acceptance(self): + """Accepted tokens per live row, bonus token included, by `spec["accept"]`.""" + mode = self.spec.get("accept", "random") + rows = self.live.numel() + if mode == "full": + return torch.full((rows,), self.t, dtype=torch.int32, device=DEV) + if mode == "one": + return torch.ones(rows, dtype=torch.int32, device=DEV) + assert mode == "random", mode + return torch.randint(1, self.t + 1, (rows,), generator=self.cpu, dtype=torch.int32).to(DEV) + + def accept(self, accepted): + _advance_replay_state(self.meta, self.live, accepted) + if self.exact: + self.x_pnat[self.live.long()] += accepted + + def writes(self): + """Where this step's calls append, by the documented rule, which slots wrap, and how + many append exactly up to the history's end: `[slots, 2, HIST]` mask, wrapping slot + ids, the `pnat + t == HIST` count.""" + s = self.live.long() + p, b = self.pnat[s].long(), self.buf[s].long() + wrap = p + self.t > HIST + half, start = torch.where(wrap, 1 - b, b), torch.where(wrap, 0, p) + mask = torch.zeros(self.slots, 2, HIST, dtype=torch.bool, device=DEV) + for t in range(self.t): + mask[s, half, start + t] = True + return mask, s[wrap], int((p + self.t == HIST).sum()) + + +def _deviation(got, want): + """The largest element deviation relative to its block's absmax (the + quantity `_compare` gates).""" + block = got.shape[-2] * got.shape[-1] + g, w = got.float().reshape(-1, block), want.float().reshape(-1, block) + keep = ~torch.isnan(w).any(1) + g, w = g[keep], w[keep] + if not w.numel(): + return 0.0 + scale = w.abs().amax(1, keepdim=True).clamp_min(1e-30) + return ((g - w).abs() / scale).max().item() + + +def _check_duplicate_padding_slot(cell) -> None: + """The last rows of a batch share one slot, as CUDA-graph padding repeats the manager's + dummy slot: the rows on unique slots, and every slot but the shared one, are held to + the reference run on those rows alone -- the padding rows change nothing about them. + The shared slot's state and the padding rows' outputs are unspecified.""" + bench = _Bench(cell.spec, seed=_seed(cell.why), exact=False) + dup = cell.spec["dup"] + unique = bench.n - len(dup) + assert dup == tuple(range(unique, bench.n)), "the padding rows trail the batch" + shared = int(bench.state_indices[unique]) + assert (bench.state_indices[unique:] == shared).all() + assert shared not in bench.live.tolist() + others = torch.ones(bench.slots, dtype=torch.bool, device=DEV) + others[shared] = False + for _ in range(bench.steps): + for layer in range(bench.layers): + act = bench.inputs() + kw = bench.kwargs(layer, act) + alone = dict( + kw, + q=kw["q"][:unique], + k=kw["k"][:unique], + v=kw["v"][:unique], + g=kw["g"][:unique], + beta=kw["beta"][:unique], + state_indices=kw["state_indices"][:unique], + packed_qkv=act["packed"][: unique * bench.t], + ) + ref = op.reference(**alone) + got = op(**kw) + pool, u, k, G = bench.state(layer) + _compare( + (got[:unique], pool[others], u[others], k[others], G[others]), + (ref[0], *(x[others] for x in ref[1:])), + ) + bench.accept(bench.draw_acceptance()) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + if "dup" in cell.spec: + return _check_duplicate_padding_slot(cell) + bench = _Bench(cell.spec, seed=_seed(cell.why)) + pad = list(cell.spec.get("pad", ())) + wraps = appends = to_the_end = 0 + for _ in range(bench.steps): + mask, wrapped, at_end = bench.writes() + wraps += wrapped.numel() + to_the_end += at_end + appends += bench.live.numel() - wrapped.numel() + bench.prepare() + after_commit = [] + for layer in range(bench.layers): + act = bench.inputs() + kw = bench.kwargs(layer, act) + pool, u, k, G = bench.state(layer) + before = dict( + pool=pool.clone(), + u=u.clone(), + k=k.clone(), + G=G.clone(), + old_beta=bench.old_dA[layer].clone(), + out_buf=None if act["out_buf"] is None else act["out_buf"].clone(), + ) + ref = op.reference(**kw) + exact = op.reference(**bench.exact_kwargs(layer, kw)) + committed = None + if bench.path == "all_layer": + committed = op.reference( + **dict(kw, replay_work_items=None, n_writes=None, use_all_layer_commit=False) + ) + got = op(**kw) + _compare((got, *bench.state(layer)), ref) + # The locks: in place, and nothing outside the documented writes. + pool, u, k, G = bench.state(layer) + assert torch.equal(bench.old_dA[layer], before["old_beta"]), "old_beta is never written" + assert torch.equal(u[~mask], before["u"][~mask]), "old_u written outside the append" + assert torch.equal(k[~mask], before["k"][~mask]), "old_k written outside the append" + Gt, Gt0 = G.transpose(-1, -2), before["G"].transpose(-1, -2) + assert torch.equal(Gt[~mask], Gt0[~mask]), "old_G written outside the append" + untouched = torch.ones(bench.slots, dtype=torch.bool, device=DEV) + if bench.path == "small": + untouched[wrapped] = False + assert torch.equal(pool[untouched], before["pool"][untouched]), ( + "the pool moved where the call commits nothing (a partitioned call commits none)" + ) + if kw["output"] is not None: + assert got.data_ptr() == kw["output"].data_ptr(), "output must be written in place" + outside = torch.ones(act["out_buf"].shape[0], dtype=torch.bool, device=DEV) + outside[act["lo"] : act["lo"] + bench.n * bench.t] = False + assert torch.equal(act["out_buf"][outside], before["out_buf"][outside]), ( + "rows of the output buffer outside the view were touched" + ) + for row in pad: + assert (got[row] == OUT_SENTINEL).all(), f"padding row {row} written" + bench.adopt_exact(layer, exact) + after_commit.append((got, exact, committed)) + bench.commit() + for layer, (got, exact, committed) in enumerate(after_commit): + if committed is not None: + # The partitioned call plus the manager's commit is the call + # that commits in-kernel. + _compare((got, *bench.state(layer)), committed) + _compare((got, None, None, None, None), (exact[0], None, None, None, None)) + bench.accept(bench.draw_acceptance()) + assert wraps > 0 and appends > 0, ( + f"the cell never wrapped or never appended ({wraps}, {appends})" + ) + if cell.spec.get("accept") == "full": + live = bench.live.numel() + assert wraps >= 3 * live and to_the_end >= 3 * live, ( + f"full acceptance should wrap every row at least 3 times and append up to the " + f"history's end as often ({wraps}, {to_the_end} over {live} rows)" + ) + + +# ── CUDA graphs ─────────────────────────────────────────────────────────── + + +_GRAPH_CASES = [(n, t, torch.float32) for t in range(2, 9) for n in (1, 8, 15, 16, 32)] + [ + (n, t, torch.bfloat16) for t in (2, 4, 8) for n in (1, 15, 16, 32) +] + + +@pytest.mark.parametrize( + "n,t,state_dtype", + _GRAPH_CASES, + ids=[f"{n}-{t}-{'fp32' if d is torch.float32 else 'bf16'}" for n, t, d in _GRAPH_CASES], +) +def test_cuda_graph_replay_on_new_data_is_bitwise_eager(n, t, state_dtype) -> None: + """Every T 2..8 (MTP draft lengths 1..7) over an fp32 checkpoint, and T 2, 4 and 8 over a + bf16 one, at the batches either side of the partition threshold and at 1, 8 and 32: two + layers captured with, from 16 sequences on, the all-layer commit after them; replayed + three times on new activations, with acceptance in between, against the same steps run + eagerly on a twin state.""" + assert ALL_LAYER_COMMIT_MIN_BATCH == CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE + path = "all_layer" if n >= ALL_LAYER_COMMIT_MIN_BATCH else "small" + spec = dict(path=path, n=n, t=t, layers=2, slots=40, state_dtype=state_dtype) + seed = 7000 + 10 * n + t + eager, graph, warm = (_Bench(spec, seed, exact=False) for _ in range(3)) + acts = [warm.inputs() for _ in range(2)] + warm.prepare() + for layer in range(2): + op(**warm.kwargs(layer, acts[layer])) + warm.commit() + torch.cuda.synchronize() + + graph.prepare() + captured = torch.cuda.CUDAGraph() + with torch.cuda.graph(captured): + outs_graph = [op(**graph.kwargs(layer, acts[layer])) for layer in range(2)] + graph.commit() + assert torch.equal(graph.backing, eager.backing) and torch.equal(graph.old_x, eager.old_x) + + refill = torch.Generator(device=DEV).manual_seed(seed + 1) + for _ in range(3): + for act in acts: + for key in ("packed", "a", "b"): + act[key].copy_( + torch.randn(act[key].shape, generator=refill, device=DEV, dtype=act[key].dtype) + ) + eager.prepare() + outs_eager = [op(**eager.kwargs(layer, acts[layer])) for layer in range(2)] + eager.commit() + graph.prepare() + captured.replay() + torch.cuda.synchronize() + for e, g in zip(outs_eager, outs_graph): + assert torch.equal(e, g), "replayed output differs from eager" + for name in ("backing", "old_x", "old_B", "old_dt", "pnat", "buf"): + assert torch.equal(getattr(eager, name), getattr(graph, name)), f"{name} differs" + accepted = eager.draw_acceptance() + eager.accept(accepted) + graph.accept(accepted) + + +# ── the harness can see what it exists to see ───────────────────────────── + + +def _drive(bench, fault): + """Run `bench` with one protocol fault injected; return the first step at + which the twin comparison rejected an output, or None.""" + for step in range(bench.steps): + bench.prepare() + outs = [] + for layer in range(bench.layers): + kw = bench.kwargs(layer, bench.inputs()) + if fault == "committed twice": + kw = dict(kw, replay_work_items=None, n_writes=None, use_all_layer_commit=False) + exact = op.reference(**bench.exact_kwargs(layer, kw)) + bench.adopt_exact(layer, exact) + outs.append((op(**kw), exact[0])) + if fault != "commit skipped": + bench.commit() + try: + for got, want in outs: + _compare((got, None, None, None, None), (want, None, None, None, None)) + except AssertionError: + return step + accepted = bench.draw_acceptance() + if fault == "rejected tokens kept": + _advance_replay_state(bench.meta, bench.live, torch.full_like(accepted, bench.t)) + bench.x_pnat[bench.live.long()] += accepted + else: + bench.accept(accepted) + return None + + +@pytest.mark.parametrize("t", (4, 8)) +@pytest.mark.parametrize( + "fault", ("none", "commit skipped", "committed twice", "rejected tokens kept") +) +def test_the_twin_sees_a_broken_protocol(fault, t) -> None: + """The twin comparison is what certifies the commit and the bookkeeping, so + it has to fail when either is broken: the partitioned path with the + manager's commit left out; a 16-row batch called without work items while + the manager still commits (the checkpoint advanced twice); and the small + path advanced as if every drafted token had been accepted. Without a fault + the same run passes.""" + path = "small" if fault == "rejected tokens kept" else "all_layer" + n = 8 if path == "small" else 16 + bench = _Bench(dict(path=path, n=n, t=t, steps=4), seed=93) + first = _drive(bench, fault) + if fault == "none": + assert first is None + else: + assert first is not None, f"{fault} passed the twin comparison at T={t}" + + +@pytest.mark.parametrize("state_dtype", (torch.float32, torch.bfloat16), ids=("fp32", "bf16")) +def test_the_gate_rejects_near_misses(state_dtype) -> None: + """`_compare` against plausible wrong answers, each the entry's own + reference driven on transformed inputs, at a call where one row wraps (and + commits) and the others append -- at each checkpoint dtype's own state gate.""" + bench = _Bench(dict(path="small", n=8, t=4, state_dtype=state_dtype), seed=91, exact=False) + kw = bench.kwargs(0, bench.inputs()) + pool, u, k, G = (x.clone() for x in bench.state(0)) + base = dict(kw, ssm_states=pool, old_u=u, old_k=k, old_G=G) + ref = op.reference(**base) + got = op(**kw) + out = (got, *bench.state(0)) + _compare(out, ref) + p, b = kw["prev_num_accepted_tokens"], kw["cache_buf_idx"] + + not_restarted = ref[4].clone() + _, wrapped, _ = bench.writes() + for s in wrapped.tolist(): + g_start = G[s, b[s].long(), :, p[s].long() - 1] + not_restarted[s, 1 - b[s].long(), :, : bench.t] += g_start[:, None] + rivals = { + "history ignored": op.reference(**dict(base, prev_num_accepted_tokens=torch.zeros_like(p))), + "the other history half": op.reference(**dict(base, cache_buf_idx=1 - b)), + "one pending token dropped": op.reference( + **dict(base, prev_num_accepted_tokens=(p - 1).clamp(min=0)) + ), + "checkpoint zeroed": op.reference(**dict(base, ssm_states=torch.zeros_like(pool))), + "G not restarted at a wrap": ref[:4] + (not_restarted,), + } + for rival in rivals.values(): + with pytest.raises(AssertionError): + _compare(out, rival) + + +# ── properties beyond one reference comparison ──────────────────────────── + + +def test_the_wrapped_kernels_carry_no_autotuner() -> None: + """Plain `triton.jit`: no trial launch can consume the in-place state on a + process's first call.""" + for kernel in (_impl._cached_replay_kernel, _impl._cached_replay_layered_commit_kernel): + names = [cls.__name__ for cls in type(kernel).__mro__] + assert "Autotuner" not in names and "Heuristics" not in names, names + + +def _clone_state(bench): + return {name: getattr(bench, name).clone() for name in ("backing", "old_x", "old_B", "old_dt")} + + +def _restore_state(bench, saved): + for name, value in saved.items(): + getattr(bench, name).copy_(value) + + +@pytest.mark.parametrize("state_dtype", (torch.float32, torch.bfloat16), ids=("fp32", "bf16")) +def test_the_two_commit_paths_are_bitwise_interchangeable(state_dtype) -> None: + """The same 16 rows through the in-kernel commit, and through the + partitioned call plus the manager's all-layer commit: bit for bit, at either + checkpoint dtype.""" + bench = _Bench( + dict(path="all_layer", n=16, t=4, layers=2, state_dtype=state_dtype), seed=95, exact=False + ) + acts = [bench.inputs() for _ in range(2)] + saved = _clone_state(bench) + small = [ + op( + **dict( + bench.kwargs(i, acts[i]), + replay_work_items=None, + n_writes=None, + use_all_layer_commit=False, + ) + ) + for i in range(2) + ] + after_small = _clone_state(bench) + _restore_state(bench, saved) + bench.prepare() + assert 0 < int(bench.n_writes) < bench.n, "want both wrapping and appending rows" + partitioned = [op(**bench.kwargs(i, acts[i])) for i in range(2)] + bench.commit() + for a, b in zip(small, partitioned): + assert torch.equal(a, b) + for name, value in after_small.items(): + assert torch.equal(getattr(bench, name), value), name + + +def test_packed_qkv_supersedes_q_k_v_and_launch_knobs_are_neutral() -> None: + """With `packed_qkv` the kernel reads q/k/v from it alone -- their own + contents are never read; the separate contiguous form, PDL on or off, and a + second run are all bitwise the same call.""" + bench = _Bench(dict(path="small", n=6, t=3, layers=1), seed=97, exact=False) + act = bench.inputs() + saved = _clone_state(bench) + + def run(**override): + _restore_state(bench, saved) + out = op(**dict(bench.kwargs(0, act), **override)).clone() + return out, _clone_state(bench) + + base = run() + q, k, v = act["q"], act["k"], act["v"] + variants = { + "separate q/k/v": run( + q=q.contiguous(), k=k.contiguous(), v=v.contiguous(), packed_qkv=None + ), + "q/k/v contents zeroed": run( + q=torch.zeros_like(q.contiguous()), + k=torch.zeros_like(k.contiguous()), + v=torch.zeros_like(v.contiguous()), + ), + "launch_with_pdl=True": run(launch_with_pdl=True), + "second run": run(), + } + for name, (out, state) in variants.items(): + assert torch.equal(out, base[0]), name + for key, value in state.items(): + assert torch.equal(value, base[1][key]), (name, key) + + +# ── guards: accepted and silently wrong ────────────────────────────────── + + +def test_a_short_state_indices_reads_past_its_end_and_is_refused() -> None: + """Unguarded, three indices for four rows: the fourth row reads whatever + follows the view -- here a slot the call never named -- and writes its + history.""" + bench = _Bench(dict(path="small", n=4, t=2, layers=1), seed=101, exact=False) + stranger = next(s for s in range(bench.slots) if s not in bench.live.tolist()) + parent = torch.cat( + [bench.state_indices[:3], torch.tensor([stranger], dtype=torch.int32, device=DEV)] + ) + kw = dict(bench.kwargs(0, bench.inputs()), state_indices=parent[:3]) + before = bench.old_x[0, stranger].clone() + op(**kw) + assert not torch.equal(bench.old_x[0, stranger], before), "the hazard this guard covers is gone" + with pytest.raises(AssertionError, match="state_indices has shape"): + with validating(op): + op(**kw) + + +def test_a_history_off_history_size_is_misread_and_is_refused() -> None: + """Unguarded, 16-deep buffers called with history_size=12: the kernel strides + them as 12 deep and reads other slots' entries. The same call on 12-deep + copies of the same history is the right answer. (12, not 8: below 9 the + kernel does not compile -- see `test_op_refuses_what_it_refuses`.)""" + bench = _Bench(dict(path="small", n=4, t=4, layers=1), seed=103, exact=False) + live = bench.live.long() + bench.pnat[live] = torch.tensor([1, 2, 3, 4], dtype=torch.int32, device=DEV) + kw = bench.kwargs(0, bench.inputs()) + shallow = dict( + kw, + old_u=kw["old_u"][:, :, :12].clone(), + old_k=kw["old_k"][:, :, :12].clone(), + old_G=kw["old_G"][..., :12].clone(), + ssm_states=kw["ssm_states"].clone(), + history_size=12, + ) + ref = op.reference(**shallow) + right = op(**shallow) + _compare((right, None, None, None, None), (ref[0], None, None, None, None)) + wrong = op(**dict(kw, history_size=12)) + gap = _deviation(wrong, right) / BF16_ULP + assert gap > 64, "the hazard this guard covers is gone" + with pytest.raises(AssertionError, match="old_u has shape"): + with validating(op): + op(**dict(kw, history_size=12)) + + +@pytest.mark.parametrize( + "name, form", + ( + ("history", "fewer slots"), + ("prev_num_accepted_tokens", "fewer slots"), + ("cache_buf_idx", "fewer slots"), + ("prev_num_accepted_tokens", "strided"), + ("cache_buf_idx", "strided"), + ), +) +def test_bookkeeping_that_does_not_cover_the_pool_is_refused(name, form) -> None: + """The kernel addresses the history buffers, `prev_num_accepted_tokens` and + `cache_buf_idx` by the pool's slot id: one with fewer slots than `ssm_states` is read and + written past its end for the slots beyond it, and a strided per-slot vector is read at the + wrong elements. Only the refusal runs; the unguarded call is out of bounds.""" + bench = _Bench(dict(path="small", n=4, t=2, layers=1), seed=109, exact=False) + kw = bench.kwargs(0, bench.inputs()) + last = bench.slots - 1 + if name == "history": + kw = dict(kw, old_u=kw["old_u"][:last], old_k=kw["old_k"][:last], old_G=kw["old_G"][:last]) + match = "history buffers have" + else: + vec = kw[name] + if form == "fewer slots": + kw[name] = vec[:last] + else: + kw[name] = torch.stack([vec, vec], 1)[:, 0] + match = f"{name} has shape" + with pytest.raises(AssertionError, match=match): + with validating(op): + op(**kw) + + +def test_op_refuses_what_it_refuses() -> None: + """Loud refusals; the guard admits each so the op's own error surfaces.""" + bench = _Bench(dict(path="all_layer", n=16, t=4, layers=1), seed=105, exact=False) + bench.prepare() + kw = bench.kwargs(0, bench.inputs()) + with pytest.raises(ValueError, match="Partitioned replay requires the all-layer commit"): + with validating(op): + op(**dict(kw, use_all_layer_commit=False)) + with pytest.raises(ValueError, match="use_all_layer_commit requires replay work items"): + with validating(op): + op(**dict(kw, replay_work_items=None, n_writes=None)) + with pytest.raises(ValueError, match="must either both be set or both be None"): + with validating(op): + op(**dict(kw, n_writes=None)) + with pytest.raises(AssertionError, match="output must be contiguous"): + wide = torch.empty(16, 4, HV, 2 * V, dtype=torch.bfloat16, device=DEV) + with validating(op): + op(**dict(kw, output=wide[..., :V])) + with pytest.raises(AssertionError, match="g head dimension must be contiguous"): + with validating(op): + op(**dict(kw, g=kw["g"].transpose(1, 2).contiguous().transpose(1, 2))) + + def shallow(depth): + return dict( + kw, + history_size=depth, + old_u=kw["old_u"][:, :, :depth].contiguous(), + old_k=kw["old_k"][:, :, :depth].contiguous(), + old_G=kw["old_G"][..., :depth].contiguous(), + ) + + with pytest.raises(AssertionError): # `history_size >= T`, a bare assert + with validating(op): + op(**shallow(2)) + # The history tile is next_pow2(history_size) wide and is a `tl.dot` + # reduction, which needs 16: history_size 8 (>= T = 4) fails to compile. + with pytest.raises(CompilationError, match="K >= 16"): + with validating(op): + op(**shallow(8)) + # Nine rows per sequence (max_draft_len 8) is past the kernel's `T <= 8`, a bare assert + # before the launch. + wide = _Bench(dict(path="small", n=2, t=9, layers=1), seed=106, exact=False) + with pytest.raises(AssertionError): + with validating(op): + op(**wide.kwargs(0, wide.inputs())) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_sigmoid_gating_delta_rule_update.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_sigmoid_gating_delta_rule_update.py index 114df72f3257..313b228ff644 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_sigmoid_gating_delta_rule_update.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_sigmoid_gating_delta_rule_update.py @@ -2,9 +2,11 @@ # SPDX-License-Identifier: Apache-2.0 """GPU test for the fused_sigmoid_gating_delta_rule_update entry. -Everything here certifies the vendored Triton branch; the fp32 recurrent pool -this target runs with keeps the FlashInfer branch out, which the first test -below proves. +Everything here certifies the vendored Triton branch. The target's two decode calls +keep the FlashInfer branch out -- the packed call over an fp32 pool, and the rows form +(no `cu_seqlens`) over a bf16 one -- which `test_certified_branch_is_the_vendored_triton_kernel` +proves. `test_rows_form_is_batch_invariant` shows the property the bf16 call takes the rows +form for. """ import zlib @@ -35,7 +37,9 @@ def _compare( Both at torch's default dtype-aware band: the kernel and the reference accumulate the same fp32 sequence and round once (bf16 `o`: measured - 3.05e-05 max abs at N=32, 461 slots; fp32 pool 2.4e-07). Five wrong + 3.05e-05 max abs at N=32, 461 slots; fp32 pool 2.4e-07). Over a bf16 pool the + state is rounded once on store, as the reference rounds it, so a stored value + can sit one bf16 step (2^-8 relative) from the reference's, inside the band. Five wrong variants (a/b swapped, norm dropped, pool `(V, K)` transposed, initial state ignored, q-head order reversed) land 1.7e-2 to 4.5 away, three orders of magnitude outside. The pool is compared whole, so the @@ -62,6 +66,16 @@ def _compare( _TAIL = 5 +def _rows_form(mixed, a, b): + """The target's bf16 decode operands: q/k/v `[N, 1, heads, 128]` views of the packed + `[N, 10240]` block, reshaped whole so the size-1 token axis carries the row stride, and + a/b `[N, 1, 48]`.""" + n = mixed.shape[0] + qkv = mixed.view(n, 1, MIXED_QKV_WIDTH // HEAD_K_DIM, HEAD_K_DIM) + q, k, v = torch.split(qkv, [NUM_K_HEADS, NUM_K_HEADS, NUM_V_HEADS], 2) + return q, k, v, a.reshape(n, 1, NUM_V_HEADS), b.reshape(n, 1, NUM_V_HEADS) + + def _build(spec, seed): """`__call__`'s 15 positional arguments, plus the backing buffers some layouts view. q/k/v are views over one packed `[T, 10240]` conv output, @@ -88,7 +102,7 @@ def _build(spec, seed): A_log = torch.randn(NUM_V_HEADS, generator=g, device=DEV) dt_bias = torch.randn(NUM_V_HEADS, generator=g, device=DEV) - extras = {} + extras = {"mixed": mixed} slots = spec["slots"] pool = None if slots is not None: @@ -110,6 +124,9 @@ def _build(spec, seed): 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 + elif layout == "rows": + q, k, v, a, b = _rows_form(mixed, a, b) + cu_seqlens = None output = None if spec.get("output") == "prefix": extras["buf"] = torch.full( @@ -167,9 +184,10 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res 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.""" + """Both target decode calls keep the FlashInfer arm out, armed both ways: its entry + point is replaced by one that raises; the packed call over an fp32 pool and the rows + form over a bf16 pool complete, and the packed call over a bf16 pool in the same + process does raise -- so the replaced entry point is known to be reachable.""" def _boom(*args, **kwargs): raise AssertionError("FlashInfer arm taken") @@ -181,6 +199,10 @@ def _boom(*args, **kwargs): 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) + rows = dict(seq_lens=(1,) * 32, slots=64, ab="separate", pool_dtype=torch.bfloat16) + args, _ = _build(dict(rows, layout="rows"), seed=102) + assert op(*args).shape == (32, 1, 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"): @@ -189,6 +211,109 @@ def _boom(*args, **kwargs): _impl._flashinfer_gdn_decode = original +#: Batch sizes either side of every FlashInfer dispatch boundary at 48 value heads (its +#: kernel and tile change at 2|3, 6|7, 10|11 and 21|22) and the engine's CUDA-graph batch +#: sizes up to 32. +_BATCHES = (1, 2, 3, 10, 11, 16, 21, 22, 32) +_FIXED_SLOT = 301 +_FIXED_ROW_SLOTS = 512 + + +def test_rows_form_is_batch_invariant() -> None: + """The property the target's bf16 decode call relies on: a row's output and state are + bitwise independent of the batch it rides in, at batch sizes either side of every + FlashInfer dispatch boundary and at both ends of the batch. One fixed row (activations + and slot 301 of a bf16 pool) sits at position 0 and N-1 of random N-row batches.""" + gen = torch.Generator(device=DEV).manual_seed(3) + cpu = torch.Generator().manual_seed(3) + pool0 = ( + torch.randn( + _FIXED_ROW_SLOTS, NUM_V_HEADS, HEAD_V_DIM, HEAD_K_DIM, generator=gen, device=DEV + ) + * 0.05 + ).to(torch.bfloat16) + row = torch.randn(MIXED_QKV_WIDTH, generator=gen, device=DEV, dtype=torch.bfloat16) + row_a, row_b = ( + torch.randn(NUM_V_HEADS, generator=gen, device=DEV, dtype=torch.bfloat16) for _ in range(2) + ) + A_log = torch.randn(NUM_V_HEADS, generator=gen, device=DEV) + dt_bias = torch.randn(NUM_V_HEADS, generator=gen, device=DEV) + first = None + for n in _BATCHES: + for pos in sorted({0, n - 1}): + g = torch.Generator(device=DEV).manual_seed(1000 * n + pos) + mixed = torch.randn(n, MIXED_QKV_WIDTH, generator=g, device=DEV, dtype=torch.bfloat16) + a = torch.randn(n, NUM_V_HEADS, generator=g, device=DEV, dtype=torch.bfloat16) + b = torch.randn(n, NUM_V_HEADS, generator=g, device=DEV, dtype=torch.bfloat16) + mixed[pos], a[pos], b[pos] = row, row_a, row_b + others = [s for s in torch.randperm(_FIXED_ROW_SLOTS, generator=cpu).tolist()] + others = [s for s in others if s != _FIXED_SLOT][:n] + others[pos] = _FIXED_SLOT + slots = torch.tensor(others, dtype=torch.int32, device=DEV) + pool = pool0.clone() + q, k, v, a_in, b_in = _rows_form(mixed, a, b) + out = op(A_log, a_in, dt_bias, 1.0, 20.0, q, k, v, b_in, pool, slots, None, True, None) + got = (out.reshape(n, NUM_V_HEADS, HEAD_V_DIM)[pos].clone(), pool[_FIXED_SLOT].clone()) + first = first or got + assert torch.equal(got[0], first[0]) and torch.equal(got[1], first[1]), ( + f"the fixed row differs at batch {n}, position {pos}" + ) + + +def test_rows_form_is_the_packed_triton_call_bitwise() -> None: + """At either pool dtype the rows form computes what the packed Triton call computes, + bit for bit -- output and pool. The packed call is forced onto the Triton kernel for + the bf16 comparison.""" + for pool_dtype in (torch.float32, torch.bfloat16): + for n in (1, 8, 32): + spec = dict(seq_lens=(1,) * n, slots=64, ab="separate", pool_dtype=pool_dtype) + packed, _ = _build(spec, seed=300 + n) + rows, _ = _build(dict(spec, layout="rows"), seed=300 + n) + with pytest.MonkeyPatch.context() as mp: + mp.setattr(_impl, "_can_use_flashinfer_gdn_decode", lambda *a, **k: False) + o_packed = op(*packed) + o_rows = op(*rows) + assert torch.equal(o_rows.reshape(o_packed.shape), o_packed), (pool_dtype, n) + assert torch.equal(rows[9], packed[9]), (pool_dtype, n) + + +@pytest.mark.parametrize("n", (1, 16, 32)) +def test_rows_form_cuda_graph_replay_is_bitwise_eager(n) -> None: + """The bf16 rows form captured at batch `n` and replayed on new activations, slots and + pool content, against the same call run eagerly. The kernel takes the batch as a + constexpr, so each batch size compiles on its eager call, before the capture.""" + spec = dict( + seq_lens=(1,) * n, + slots=ENGINE_SLOTS, + ab="separate", + layout="rows", + pool_dtype=torch.bfloat16, + ) + args, extras = _build(spec, seed=400 + n) + static = list(args) + for _ in range(2): + op(*static) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out_static = op(*static) + for trial in range(3): + slots = tuple(7 * trial + 13 * i for i in range(n)) + fresh, fresh_extras = _build(dict(spec, indices=slots), seed=500 + 10 * n + trial) + for position in (0, 1, 2, 8, 9, 10): + static[position].copy_(fresh[position]) + extras["mixed"].copy_(fresh_extras["mixed"]) + pool_before = static[9].clone() + graph.replay() + torch.cuda.synchronize() + replayed, replayed_pool = out_static.clone(), static[9].clone() + eager_args = list(fresh) + eager_args[9] = pool_before + eager = op(*eager_args) + assert torch.equal(replayed, eager), (n, trial) + assert torch.equal(replayed_pool, pool_before), (n, trial) + + 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.""" @@ -226,6 +351,35 @@ def test_guard_refuses_a_noncontiguous_output_buffer() -> None: op(*guarded) +def test_guard_refuses_rows_views_whose_token_stride_is_not_the_row_stride() -> None: + """Unguarded, q/k/v reshaped to `[N, 1, heads, 128]` from column slices of the packed + block carry `stride(1) = heads * 128`: without `cu_seqlens` the kernel reads row n at + `n * stride(1)`, from other rows' columns.""" + n, kd = 8, NUM_K_HEADS * HEAD_K_DIM + spec = dict( + seq_lens=(1,) * n, slots=64, ab="separate", layout="rows", pool_dtype=torch.bfloat16 + ) + args, extras = _build(spec, seed=600) + mixed = extras["mixed"] + pool_before = args[9].clone() + ref = op.reference(*args) + bad = list(args) + bad[5] = mixed[:, :kd].reshape(n, 1, NUM_K_HEADS, HEAD_K_DIM) + bad[6] = mixed[:, kd : 2 * kd].reshape(n, 1, NUM_K_HEADS, HEAD_K_DIM) + bad[7] = mixed[:, 2 * kd :].reshape(n, 1, NUM_V_HEADS, HEAD_V_DIM) + assert all(bad[i].stride(1) != bad[i].stride(0) for i in (5, 6, 7)) + for i in (5, 6, 7): + assert torch.equal(bad[i], args[i]), "the same values, other strides" + got = op(*bad) + with pytest.raises(AssertionError): + _compare((got, bad[9]), ref) + guarded = list(bad) + guarded[9] = pool_before + with pytest.raises(AssertionError, match=r"stride\(0\) must equal"): + with validating(op): + op(*guarded) + + 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) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_gdn_replay_verify.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_gdn_replay_verify.py new file mode 100644 index 000000000000..239bcb1c13f8 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_gdn_replay_verify.py @@ -0,0 +1,364 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the target's Gated-DeltaNet verify step in both recurrent modes. + +Drives the two targets' own `_gdn_verify` -- the prefill target's (generation rows +behind a context prefix) and the decode target's -- over consecutive verify steps of +this checkpoint's GDN geometry, with the cache manager's bookkeeping run in between as +`update_mamba_states` runs it: the conv window promotion in both modes, then either the +legacy per-step state promotion, or the cached replay's all-layer commit (partitioned +batches only) and `_advance_replay_state`. + +Three things are checked at every step, with a random partial acceptance (or, in the +full-acceptance cases, every draft accepted): + +* the replay mode's output against the legacy mode's on identical inputs and identical + accepted histories: the two modes compute one recurrence with different kernels and + state disciplines, so a mis-bound history buffer, a wrong partition condition or a + skipped commit drifts off within a step; +* the prefill and the decode target's replay calls bitwise equal to each other (each on + its own copy of the state); +* the conv pool identical across the modes (it is promoted the same way). +""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 + modeling as M, +) +from tensorrt_llm._torch.modules.fla.cached_replay import ( + CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE, + commit_gdn_cached_replay_history_layers, +) +from tensorrt_llm._torch.modules.mamba.mamba2_metadata import _build_replay_work_items_torch +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( + ReplayStateUpdateMetadata, + _advance_replay_state, +) + +assert torch.cuda.is_available(), "the GDN verify test requires a CUDA device" + +DEV = "cuda" +H, HV, K, V = 16, 48, 128, 128 +KEY_DIM, VALUE_DIM = H * K, HV * V +CONV_DIM = 2 * KEY_DIM + VALUE_DIM +CONV_KERNEL = 4 +HIST = 16 +LAYERS = 2 +SLOTS = 40 +STEPS = 7 + + +def _steps(t): + """Enough steps that every slot fills its history at least once: each step accepts + at least one token, and a slot wraps once `pnat + t > HIST`.""" + return HIST - t + 3 + + +BF16_ULP = 2.0**-8 +#: Replay against legacy, per token block: the replay kernel rounds the history's update +#: vectors and keys to bf16 and rebuilds the state from them, where the legacy kernel +#: carries an fp32 state, so the two agree to a few bf16 ulps of the block. +ELEMENT_GATE = 6 * BF16_ULP +RMS_GATE = 1.5 * BF16_ULP + + +class _Mode: + """One cache manager's worth of state for `LAYERS` GDN layers, in one mode, plus the + `core` / `L` stand-ins `_gdn_verify` reads.""" + + def __init__(self, replay, t, conv0, ssm0, weights, gating, max_batch): + self.conv = [c.clone() for c in conv0] + self.int_conv = [ + torch.zeros(max_batch, t, CONV_DIM, CONV_KERNEL - 1, dtype=torch.bfloat16, device=DEV) + for _ in range(LAYERS) + ] + # V2 lays the layers out inside each slot; the all-layer commit is handed that + # affine layout. + self.backing = torch.stack(ssm0, 1).contiguous() + self.ssm = [self.backing[:, i] for i in range(LAYERS)] + bound = dict( + conv=dict(enumerate(self.conv)), + int_conv=dict(enumerate(self.int_conv)), + int_ssm=None, + replay=None, + ) + if replay: + self.old_x = torch.zeros( + LAYERS, SLOTS, 2, HIST, HV, V, dtype=torch.bfloat16, device=DEV + ) + self.old_B = torch.zeros(LAYERS, SLOTS, 2, HIST, H, K, dtype=torch.bfloat16, device=DEV) + self.old_dt = torch.zeros(LAYERS, SLOTS, 2, HV, HIST, device=DEV) + self.old_dA = torch.zeros(LAYERS, SLOTS, 2, HV, HIST, device=DEV) + self.pnat = torch.zeros(SLOTS, dtype=torch.int32, device=DEV) + self.buf = torch.zeros(SLOTS, dtype=torch.int32, device=DEV) + bound["replay"] = M._Replay( + old={ + i: (self.old_x[i], self.old_B[i], self.old_dt[i], self.old_dA[i]) + for i in range(LAYERS) + }, + cache_buf_idx=self.buf, + prev_num_accepted=self.pnat, + history=HIST, + width=t, + all_layer_commit=True, + ) + self.meta = ReplayStateUpdateMetadata(self.pnat, self.buf, t, HIST) + self.mm = SimpleNamespace( + replay_work_items=torch.zeros(max_batch, 4, dtype=torch.int32, device=DEV), + replay_n_writes=torch.zeros(1, dtype=torch.int32, device=DEV), + ) + else: + self.int_ssm = [torch.zeros(max_batch, t, HV, V, K, device=DEV) for _ in range(LAYERS)] + bound["int_ssm"] = dict(enumerate(self.int_ssm)) + self.core = SimpleNamespace( + conv_dim=CONV_DIM, + key_dim=KEY_DIM, + value_dim=VALUE_DIM, + num_k_heads=H, + head_k_dim=K, + num_v_heads=HV, + head_v_dim=V, + _bound=SimpleNamespace(**bound), + ) + self.layers = [ + SimpleNamespace(conv_w=weights[i], a_log=gating[i][0], dt_bias=gating[i][1]) + for i in range(LAYERS) + ] + + def work(self, slots, n, t): + """What the forward hands `_gdn_verify` as `replay`: the metadata's prepare builds + the work items, then `_replay_work` applies the manager's partition condition.""" + partitioned = n >= CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE + if partitioned: + _build_replay_work_items_torch( + slots, + self.pnat, + self.buf, + self.mm.replay_work_items, + self.mm.replay_n_writes, + t, + HIST, + ) + replay = M._replay_work(self.core, self.mm, n) + assert replay[2] == partitioned + return replay + + def verify(self, target_cls, layer, mq, a, b, slots, int_idx, n, t, replay): + return target_cls._gdn_verify( + SimpleNamespace(core=self.core), + mq, + a, + b, + self.ssm[layer], + layer, + self.layers[layer], + slots, + int_idx, + n, + t, + replay, + ) + + def commit(self, slots, n): + """The manager's all-layer commit, then the replay bookkeeping advance is the + caller's (it needs the accepted counts).""" + commit_gdn_cached_replay_history_layers( + ssm_states=self.ssm[0], + ssm_state_num_layers=LAYERS, + ssm_state_layer_stride=self.backing.stride(1), + ssm_state_slot_stride=self.backing.stride(0), + old_u=self.old_x, + old_k=self.old_B, + old_G=self.old_dt, + replay_work_items=self.mm.replay_work_items[:n], + n_writes=self.mm.replay_n_writes, + history_size=HIST, + ) + + +def _rows(gen, rows): + """One layer's verify-row activations: the pre-conv `[q|k|v]` projection and the raw + `a` and `b` projections, each its own contiguous GEMM output.""" + return tuple( + torch.randn(rows, width, generator=gen, device=DEV).to(torch.bfloat16) + for width in (CONV_DIM, HV, HV) + ) + + +def _behind_prefix(x, prefix=5): + """`x` as the trailing rows of a buffer that starts with `prefix` context rows.""" + buf = torch.zeros(prefix + x.shape[0], x.shape[1], dtype=x.dtype, device=DEV) + buf[prefix:] = x + return torch.split(buf, [prefix, x.shape[0]], 0)[1] + + +def _promote(pool, inter, slots, accepted): + """`update_mamba_states`' copy: step `accepted - 1` of each batch row's per-step + buffer into the row's slot.""" + rows = torch.arange(slots.shape[0], device=DEV) + pool[slots.long()] = inter[rows, accepted.long() - 1] + + +def _gate(got, want, what): + block = got.shape[-2] * got.shape[-1] + g, w = got.float().reshape(-1, block), want.float().reshape(-1, block) + scale = w.abs().amax(1, keepdim=True).clamp_min(1e-30) + worst = ((g - w).abs() / scale).max().item() + rms = ((g - w).pow(2).mean().sqrt() / w.pow(2).mean().sqrt()).item() + assert worst <= ELEMENT_GATE and rms <= RMS_GATE, ( + f"{what}: replay vs legacy {worst / BF16_ULP:.2f} ulps of the block, " + f"RMS {rms / BF16_ULP:.2f} ulps" + ) + + +def _state(gen, cpu, n, t, gating=None): + """Random weights, gating parameters and pre-step pools shared by the modes, and a + random slot assignment of `n` requests.""" + slots = torch.randperm(SLOTS, generator=cpu)[:n].to(DEV, torch.int32) + weights = [ + (0.3 * torch.randn(CONV_DIM, CONV_KERNEL, generator=gen, device=DEV)).to(torch.bfloat16) + for _ in range(LAYERS) + ] + if gating is None: + gating = [ + ( + (torch.randn(HV, generator=gen, device=DEV) * 1.5 + 1.0).clamp(-7.0, 5.0), + (torch.randn(HV, generator=gen, device=DEV) * 1.7 - 4.4).clamp(-11.0, 2.5), + ) + for _ in range(LAYERS) + ] + conv0 = [ + torch.randn(SLOTS, CONV_DIM, CONV_KERNEL - 1, generator=gen, device=DEV).to(torch.bfloat16) + for _ in range(LAYERS) + ] + ssm0 = [0.3 * torch.randn(SLOTS, HV, V, K, generator=gen, device=DEV) for _ in range(LAYERS)] + return slots, weights, gating, conv0, ssm0 + + +@pytest.mark.parametrize( + "n,t,accept", + [ + (1, 2, "random"), + (8, 4, "random"), + (CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE, 4, "random"), + (20, 2, "random"), + (18, 3, "random"), + # max_draft_len 5..7: T = 6..8, the replay kernel's ceiling. + (1, 8, "random"), + (8, 8, "random"), + (CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE, 8, "random"), + (32, 8, "random"), + (20, 6, "random"), + # Every draft accepted: from a fresh slot at T = 8 the second step appends exactly up + # to the history's end (`pnat + T == 16`) and every second step after it wraps -- + # in-kernel below the partition threshold, through the all-layer commit at and above it. + (4, 8, "full"), + (24, 8, "full"), + ], + ids=lambda v: str(v), +) +def test_replay_verify_tracks_legacy_in_both_targets(n, t, accept): + gen = torch.Generator(device=DEV).manual_seed(1000 * n + t) + cpu = torch.Generator().manual_seed(1000 * n + t) + slots, weights, gating, conv0, ssm0 = _state(gen, cpu, n, t) + int_idx = torch.arange(n, dtype=torch.int32, device=DEV) + legacy = _Mode(False, t, conv0, ssm0, weights, gating, n) + prefill = _Mode(True, t, conv0, ssm0, weights, gating, n) + decode = _Mode(True, t, conv0, ssm0, weights, gating, n) + partitioned = n >= CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE + + wraps = to_the_end = 0 + for step in range(_steps(t)): + pending = prefill.pnat[slots.long()] + wraps += int((pending + t > HIST).sum()) + to_the_end += int((pending + t == HIST).sum()) + replays = {id(mode): mode.work(slots, n, t) for mode in (prefill, decode)} + for layer in range(LAYERS): + mq, a, b = _rows(gen, n * t) + # The prefill target's generation rows sit behind a context prefix: row views + # at an offset into the batch's projections. + mq_p, a_p, b_p = (_behind_prefix(x) for x in (mq, a, b)) + want = legacy.verify(M.DecodeTarget, layer, mq, a, b, slots, int_idx, n, t, None) + got_p = prefill.verify( + M.PrefillTarget, layer, mq_p, a_p, b_p, slots, int_idx, n, t, replays[id(prefill)] + ) + got_d = decode.verify( + M.DecodeTarget, layer, mq, a, b, slots, int_idx, n, t, replays[id(decode)] + ) + assert got_p.shape == got_d.shape == want.shape == (1, n * t, HV, V) + assert torch.equal(got_p, got_d), ( + f"step {step} layer {layer}: the prefill and decode targets' replay calls disagree" + ) + _gate(got_d, want, f"step {step} layer {layer}") + if accept == "full": + accepted = torch.full((n,), t, dtype=torch.int32, device=DEV) + else: + accepted = torch.randint(1, t + 1, (n,), generator=cpu, dtype=torch.int32).to(DEV) + for layer in range(LAYERS): + for mode in (legacy, prefill, decode): + _promote(mode.conv[layer], mode.int_conv[layer], slots, accepted) + _promote(legacy.ssm[layer], legacy.int_ssm[layer], slots, accepted) + for mode in (prefill, decode): + if partitioned: + mode.commit(slots, n) + _advance_replay_state(mode.meta, slots, accepted) + for layer in range(LAYERS): + assert torch.equal(prefill.conv[layer], legacy.conv[layer]) + assert torch.equal(decode.conv[layer], legacy.conv[layer]) + assert wraps > 0, "the steps never filled a history, so the checkpoint never advanced" + if accept == "full": + assert wraps >= 3 * n and to_the_end >= 3 * n, ( + f"full acceptance at T={t} should wrap every row at least 3 times and append up to " + f"the history's end as often ({wraps}, {to_the_end})" + ) + + +def test_a_wrong_partition_drifts_off_legacy(): + """The control: the work items withheld from a partitioned step while the manager + still runs its all-layer commit -- the checkpoint committed twice -- must fail the + same gate.""" + n, t = CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE, 4 + gen = torch.Generator(device=DEV).manual_seed(7) + cpu = torch.Generator().manual_seed(7) + gating = [ + (torch.full((HV,), 0.5, device=DEV), torch.full((HV,), -4.0, device=DEV)) + for _ in range(LAYERS) + ] + slots, weights, gating, conv0, ssm0 = _state(gen, cpu, n, t, gating) + int_idx = torch.arange(n, dtype=torch.int32, device=DEV) + legacy = _Mode(False, t, conv0, ssm0, weights, gating, n) + wrong = _Mode(True, t, conv0, ssm0, weights, gating, n) + failed = False + for _ in range(STEPS): + # The manager's prepare builds the work items; the step does not receive them. + _build_replay_work_items_torch( + slots, + wrong.pnat, + wrong.buf, + wrong.mm.replay_work_items, + wrong.mm.replay_n_writes, + t, + HIST, + ) + for layer in range(LAYERS): + mq, a, b = _rows(gen, n * t) + want = legacy.verify(M.DecodeTarget, layer, mq, a, b, slots, int_idx, n, t, None) + got = wrong.verify( + M.DecodeTarget, layer, mq, a, b, slots, int_idx, n, t, (None, None, False) + ) + try: + _gate(got, want, "control") + except AssertionError: + failed = True + accepted = torch.randint(1, t + 1, (n,), generator=cpu, dtype=torch.int32).to(DEV) + for layer in range(LAYERS): + _promote(legacy.conv[layer], legacy.int_conv[layer], slots, accepted) + _promote(wrong.conv[layer], wrong.int_conv[layer], slots, accepted) + _promote(legacy.ssm[layer], legacy.int_ssm[layer], slots, accepted) + wrong.commit(slots, n) + _advance_replay_state(wrong.meta, slots, accepted) + assert failed, "a doubled checkpoint commit passed the replay-vs-legacy gate" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_index_share.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_index_share.py new file mode 100644 index 000000000000..c1c91f33bf60 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_index_share.py @@ -0,0 +1,202 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU tests for the target's MTP IndexShare plumbing. + +The target's canonical selection, the capture of each request's draft-extend row +into the shared buffers (decode and mixed-batch forms), and the reuse entry's +order preservation over canonical captures. The targets are built without weights: +the capture methods read only their arguments, the metadata and the core's +draft binding. +""" + +from types import SimpleNamespace + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 + modeling as m, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_canonical_selection import ( # noqa: E501 + qsa_canonical_selection as canonical, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_mtp_reuse_topk import ( # noqa: E501 + qsa_mtp_reuse_topk as reuse, +) + +assert torch.cuda.is_available(), "the IndexShare plumbing runs on a CUDA device" + +DEV = "cuda" +K = 512 +SLOT = 1 + + +def _target(cls): + target = object.__new__(cls) + target.core = SimpleNamespace(_draft_bound=SimpleNamespace(share_slot=SLOT)) + return target + + +def _metadata(batch: int, accepted, seq_lens=None, num_contexts: int = 0): + return SimpleNamespace( + num_seqs=batch, + num_contexts=num_contexts, + seq_lens_cuda=None + if seq_lens is None + else torch.tensor(seq_lens, dtype=torch.int32, device=DEV), + qsa_mtp_num_accepted=torch.tensor(accepted, dtype=torch.int32, device=DEV), + qsa_shared_topk_indices=torch.full((3, 8, K), -7, dtype=torch.int32, device=DEV), + qsa_shared_topk_visible_blocks=torch.full((3, 8), -7, dtype=torch.int32, device=DEV), + qsa_shared_topk_captured_slots=set(), + ) + + +COLS = 2080 + + +def _scores_and_windows(rows: int, gen: torch.Generator): + """Relu-shaped scores (half zeros: ties at the budget's last score whenever fewer than + K are positive) over windows of 0 to 3K groups; the columns past a window are NaN.""" + ends = torch.randint(0, COLS + 1, (rows,), generator=gen).to(torch.int32).to(DEV) + scores = torch.relu(torch.randn(rows, COLS, generator=gen)).to(DEV) + scores = scores * (torch.rand(rows, COLS, generator=gen).to(DEV) < 0.3) + outside = torch.arange(COLS, device=DEV)[None, :] >= ends[:, None] + scores.masked_fill_(outside, float("nan")) + return scores, torch.zeros(rows, dtype=torch.int32, device=DEV), ends + + +def _select(scores, starts, ends) -> torch.Tensor: + out = torch.empty((scores.shape[0], K), dtype=torch.int32, device=DEV) + m._select(SimpleNamespace(block_topk=K), scores, starts, ends, out) + return out + + +def test_select_is_the_canonical_row_ascending_with_padding_last(): + gen = torch.Generator().manual_seed(3) + scores, starts, ends = _scores_and_windows(64, gen) + + out = _select(scores, starts, ends) + + assert out.dtype == torch.int32 + assert torch.equal(out, canonical.reference(scores, starts, ends, out)) + for row, end in zip(out.cpu().tolist(), ends.tolist()): + ids = [i for i in row if i >= 0] + assert ids == sorted(ids) and len(set(ids)) == len(ids) == min(end, K) + assert row == ids + [-1] * (K - len(ids)), "padding must stay last" + + +def test_decode_capture_takes_each_request_last_accepted_row(): + p, batch = 4, 3 + accepted = [1, 4, 2] + topk = torch.arange(batch * p * K, dtype=torch.int32, device=DEV).reshape(batch * p, K) + visible = torch.arange(100, 100 + batch * p, dtype=torch.int32, device=DEV) + md = _metadata(batch, accepted) + + _target(m.DecodeTarget)._capture_selection(md, topk, visible, batch, p) + + rows = [s * p + a - 1 for s, a in enumerate(accepted)] + torch.testing.assert_close(md.qsa_shared_topk_indices[SLOT, :batch], topk[rows]) + torch.testing.assert_close(md.qsa_shared_topk_visible_blocks[SLOT, :batch], visible[rows]) + # Other slots and rows past the batch are untouched. + assert torch.all(md.qsa_shared_topk_indices[SLOT, batch:] == -7) + assert torch.all(md.qsa_shared_topk_indices[0] == -7) + assert md.qsa_shared_topk_captured_slots == {SLOT} + + +def test_prefill_capture_takes_chunk_last_row_and_last_accepted_row(): + """Two context requests (chunks of 5 and 3 rows) then two generation + requests (4 rows each, 2 and 4 accepted). The context requests' accepted + count is whatever the worker left there and must not matter.""" + seq_lens = [5, 3, 4, 4] + accepted = [3, 1, 2, 4] + rows_total = sum(seq_lens) + topk = torch.arange(rows_total * K, dtype=torch.int32, device=DEV).reshape(rows_total, K) + visible = torch.arange(200, 200 + rows_total, dtype=torch.int32, device=DEV) + md = _metadata(4, accepted, seq_lens=seq_lens, num_contexts=2) + + _target(m.PrefillTarget)._capture_selection(md, topk, visible) + + rows = [4, 7, 8 + 1, 12 + 3] + torch.testing.assert_close(md.qsa_shared_topk_indices[SLOT, :4], topk[rows]) + torch.testing.assert_close(md.qsa_shared_topk_visible_blocks[SLOT, :4], visible[rows]) + assert md.qsa_shared_topk_captured_slots == {SLOT} + + +def test_decode_capture_at_draft_seven_takes_row_s_times_8_plus_a_minus_1(): + """max_draft_len 7: eight rows per request, accepted 1..8, so the captured row is + `s * 8 + a - 1` -- the first and the last row of a block included.""" + p, batch = 8, 4 + accepted = [1, 8, 5, 3] + topk = torch.arange(batch * p * K, dtype=torch.int32, device=DEV).reshape(batch * p, K) + visible = torch.arange(300, 300 + batch * p, dtype=torch.int32, device=DEV) + md = _metadata(batch, accepted) + + _target(m.DecodeTarget)._capture_selection(md, topk, visible, batch, p) + + rows = [s * p + a - 1 for s, a in enumerate(accepted)] + assert rows == [0, 15, 20, 26] + torch.testing.assert_close(md.qsa_shared_topk_indices[SLOT, :batch], topk[rows]) + torch.testing.assert_close(md.qsa_shared_topk_visible_blocks[SLOT, :batch], visible[rows]) + assert torch.all(md.qsa_shared_topk_indices[SLOT, batch:] == -7) + assert md.qsa_shared_topk_captured_slots == {SLOT} + + +def test_prefill_capture_at_draft_seven(): + """A mixed batch at max_draft_len 7: one 5-row context chunk, then three generation + requests of 8 rows with 8, 1 and 6 accepted.""" + seq_lens = [5, 8, 8, 8] + accepted = [2, 8, 1, 6] + rows_total = sum(seq_lens) + topk = torch.arange(rows_total * K, dtype=torch.int32, device=DEV).reshape(rows_total, K) + visible = torch.arange(400, 400 + rows_total, dtype=torch.int32, device=DEV) + md = _metadata(4, accepted, seq_lens=seq_lens, num_contexts=1) + + _target(m.PrefillTarget)._capture_selection(md, topk, visible) + + rows = [4, 5 + 7, 13 + 0, 21 + 5] + torch.testing.assert_close(md.qsa_shared_topk_indices[SLOT, :4], topk[rows]) + torch.testing.assert_close(md.qsa_shared_topk_visible_blocks[SLOT, :4], visible[rows]) + assert md.qsa_shared_topk_captured_slots == {SLOT} + + +def _reuse_case(fresh: int): + """A canonical capture of 16 requests and the reuse call's arguments for it.""" + gen = torch.Generator().manual_seed(11 + fresh) + batch = 16 + scores, starts, ends = _scores_and_windows(batch, gen) + captured_visible = ends.cpu() + captured = _select(scores, starts, ends) + shared = torch.full((8, 32, K), -1, dtype=torch.int32, device=DEV) + shared[SLOT, :batch] = captured + shared_visible = torch.zeros((8, 32), dtype=torch.int32, device=DEV) + shared_visible[SLOT, :batch] = captured_visible.to(DEV) + args = dict( + shared_indices=shared[SLOT], + shared_visible_blocks=shared_visible[SLOT], + request_indices=torch.arange(batch, dtype=torch.int32, device=DEV), + visible_blocks=(captured_visible + fresh).to(DEV), + output=torch.empty((batch, K), dtype=torch.int32, device=DEV), + ) + return args, captured_visible + + +def test_reuse_of_a_canonical_capture_stays_ascending(): + """A canonical capture plus the groups completed since is still ascending: appended ids + are newer than every captured one, and a full row gives up its oldest ids to them, so + the newest captured groups -- the ones next to the new ones -- are kept.""" + for fresh in (0, 1, 2): + args, captured_visible = _reuse_case(fresh) + + out = reuse(**args) + + assert torch.equal(out, reuse.reference(**args)) + for row, shared_row, vis in zip( + out.cpu().tolist(), args["shared_indices"].cpu().tolist(), captured_visible.tolist() + ): + ids = [i for i in row if i >= 0] + assert row == ids + [-1] * (K - len(ids)), "padding must stay last" + assert ids == sorted(ids) and len(set(ids)) == len(ids) + assert len(ids) == min(vis + fresh, K) + assert all(i < vis + fresh for i in ids) + captured = [i for i in shared_row if i >= 0] + kept = captured[max(0, len(captured) + fresh - K) :] + assert ids == kept + list(range(vis, vis + fresh)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_kv_range_probe.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_kv_range_probe.py new file mode 100644 index 000000000000..57432c1b1199 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_kv_range_probe.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The FP8 KV range probe (`TRTLLM_MODELING_V2_VALIDATE` over an fp8 pool). + +The pool stores K/V at the unit scale, so a magnitude past the e4m3 maximum saturates. The probe +reports the largest magnitude of the first eager forward's attention layers, once, and never reads +a tensor under CUDA graph capture. +""" + +from unittest.mock import Mock + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 + modeling as m, +) + +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" +LAST = 11 + + +def _kv(peak: float = 1.0): + """K and V of one layer, 5 rows x 2 heads x 8 dims, with `peak` as V's largest value.""" + k = torch.linspace(-1.0, 1.0, 80, device=DEVICE).reshape(5, 2, 8).to(torch.bfloat16) + v = torch.linspace(-0.5, 0.5, 80, device=DEVICE).reshape(5, 2, 8).to(torch.bfloat16) + v[3, 1, 4] = peak + return k, v + + +def _run_forward(probe, peaks): + """One forward's attention layers 3, 7 and LAST; returns each `observe` result.""" + return [probe.observe(layer, *_kv(peak)) for layer, peak in zip((3, 7, LAST), peaks)] + + +def test_a_forward_within_the_e4m3_range_is_silent(monkeypatch): + logger = Mock() + monkeypatch.setattr(m, "logger", logger) + + assert _run_forward(m._KvRangeProbe(LAST), (1.0, 300.0, 448.0)) == [True, True, False] + + logger.warning.assert_not_called() + + +@pytest.mark.parametrize("peak", [450.0, 1000.0, float("inf"), float("nan")]) +def test_a_magnitude_past_the_e4m3_maximum_is_reported_once(monkeypatch, peak): + logger = Mock() + monkeypatch.setattr(m, "logger", logger) + + # The oversized value is in a middle layer: the report waits for the forward's last one. + assert _run_forward(m._KvRangeProbe(LAST), (1.0, peak, 1.0)) == [True, True, False] + + logger.warning.assert_called_once() + assert "e4m3 maximum 448" in logger.warning.call_args.args[0] + + +def test_a_negative_value_counts_by_magnitude(monkeypatch): + logger = Mock() + monkeypatch.setattr(m, "logger", logger) + + assert _run_forward(m._KvRangeProbe(LAST), (1.0, -600.0, 1.0)) == [True, True, False] + + logger.warning.assert_called_once() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_observing_under_capture_reads_nothing(monkeypatch): + logger = Mock() + monkeypatch.setattr(m, "logger", logger) + probe = m._KvRangeProbe(LAST) + k, v = _kv(1000.0) + scratch = torch.zeros(1, device="cuda") + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + scratch.add_(1) # warm the allocator off the capture stream + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + kept = [] + with torch.cuda.graph(graph): + scratch.add_(1) + # The last layer, with an out-of-range value: a read here would fail the capture + # (a host sync) and a report would fire. + kept.append(probe.observe(LAST, k, v)) + + assert kept == [True], "a captured forward does not end the probe" + logger.warning.assert_not_called() + # The first eager forward after capture is what reports. + assert _run_forward(probe, (1.0, 1.0, 1000.0)) == [True, True, False] + logger.warning.assert_called_once() diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_spec_rollback.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_spec_rollback.py new file mode 100644 index 000000000000..389a0e1efb60 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_spec_rollback.py @@ -0,0 +1,435 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the target's speculative rollback handler (`SpeculativeRollback`). + +A verify step is simulated with the target's own state-writing ops -- the n-gram id +history, the dilated short conv and the QSA index store, plus an anchor-slot overwrite +standing in for the compression -- run over every draft row the way the forward runs +them, with the handler staged at the forward's call sites and committed afterwards. + +The reference shares no code with the commit: it runs the same ops from the same +pre-step state over each generation request's accepted prefix only, which is by +definition the state a verify step must leave behind. Every pool and plane is compared +bitwise in full, so a write to a slot the step does not own fails too. The step is also +run sync-free and under CUDA-graph capture and replay with the accepted counts changed +in place. +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.modeling import ( # noqa: E501 + SpeculativeRollback, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_index_store import ( # noqa: E501 + qsa_paged_index_store, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.ngram_token_history import ( # noqa: E501 + ngram_token_history, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.dilated_short_conv_paged import ( # noqa: E501 + dilated_short_conv_paged, +) + +assert torch.cuda.is_available(), "the rollback handler test requires a CUDA device" + +DEV = "cuda" +EOS = 248044 +VOCAB = 248320 +NGRAM = 3 +KERNEL = 4 +DILATION = 3 +STATE_LEN = (KERNEL - 1) * DILATION +CHANNELS = 64 +TPB = 32 +INDEX_DIM = 128 +LAYERS = 3 +COMPRESS = 4 +SLOTS = 24 +PAGES = 48 +BLOCKS_PER_SEQ = 4 + + +class _Batch: + """One verify step: `ctx_lens` context requests (fresh) followed by + generation requests of `p` rows each, at `starts` cached tokens, in the + recurrent `slots` given (contexts first). Generation request `s` reuses the + block-table row of request `dup_of[s]` when given -- the CUDA-graph padding + dummy, repeated. With `shared_page` the live index and position planes are strided + views of one per-slot buffer, as an NVFP4 KV cache's manager serves them.""" + + def __init__(self, seed, ctx_lens, gen_starts, p, slots, dup_of=None, shared_page=False): + g = torch.Generator().manual_seed(seed) + self.ctx_lens = list(ctx_lens) + self.num_contexts = len(ctx_lens) + self.n = len(gen_starts) + self.p = p + self.slots = torch.tensor(slots, dtype=torch.int32, device=DEV) + assert self.slots.shape[0] == self.num_contexts + self.n + lengths = self.ctx_lens + [p] * self.n + starts = [0] * self.num_contexts + list(gen_starts) + # Distinct pages per request; a duplicate takes its original's row. + perm = torch.randperm(PAGES, generator=g) + table = torch.reshape(perm[: len(lengths) * BLOCKS_PER_SEQ], [len(lengths), -1]) + for dup, original in (dup_of or {}).items(): + table[self.num_contexts + dup] = table[self.num_contexts + original] + self.block_table = table.to(DEV, torch.int32) + req, logical = [], [] + for r, (length, start) in enumerate(zip(lengths, starts)): + req += [r] * length + logical += range(start, start + length) + self.rows = len(req) + self.prefix = sum(self.ctx_lens) + self.req_idx = torch.tensor(req, dtype=torch.int32, device=DEV) + self.logical = torch.tensor(logical, dtype=torch.int64, device=DEV) + self.carry = torch.tensor( + [False] * self.num_contexts + [True] * self.n, dtype=torch.bool, device=DEV + ) + self.lengths = lengths + self.qsl = _qsl(lengths) + + # Pre-step state: every pool and plane random, so a stray write shows. + self.ngram0 = torch.randint(0, VOCAB, (SLOTS, NGRAM - 1), generator=g).to(DEV) + self.conv0 = torch.randn(SLOTS, CHANNELS, STATE_LEN, generator=g).to(DEV, torch.bfloat16) + self.index0 = [ + torch.randn(PAGES, TPB, 1, INDEX_DIM, generator=g).to(DEV, torch.bfloat16) + for _ in range(LAYERS) + ] + self.position0 = torch.randint(0, 1 << 20, (PAGES, TPB, 3), generator=g).to( + DEV, torch.int32 + ) + self.weight = (0.3 * torch.randn(CHANNELS, 1, KERNEL, generator=g)).to(DEV, torch.bfloat16) + + # Step inputs. + self.ids = torch.randint(0, VOCAB, (self.rows,), generator=g).to(DEV, torch.int32) + self.x = torch.randn(self.rows, CHANNELS, generator=g).to(DEV, torch.bfloat16) + self.keys = [ + torch.randn(self.rows, 1, INDEX_DIM, generator=g).to(DEV, torch.bfloat16) + for _ in range(LAYERS) + ] + self.anchor_keys = [ + torch.randn(self.rows, 1, INDEX_DIM, generator=g).to(DEV, torch.bfloat16) + for _ in range(LAYERS) + ] + self.coords = torch.randint(0, 1 << 20, (self.rows, 3), generator=g).to(DEV, torch.int32) + + # Live state the step mutates. + self.ngram = self.ngram0.clone() + self.conv = self.conv0.clone() + if shared_page: + self.index, self.position = _shared_planes(self.index0, self.position0) + else: + self.index = [t.clone() for t in self.index0] + self.position = self.position0.clone() + self.accepted = torch.ones(self.num_contexts + self.n, dtype=torch.int32, device=DEV) + + def reset(self): + self.ngram.copy_(self.ngram0) + self.conv.copy_(self.conv0) + for live, pre in zip(self.index, self.index0): + live.copy_(pre) + self.position.copy_(self.position0) + + def set_accepted(self, gen_accepted): + assert len(gen_accepted) == self.n + assert all(1 <= a <= self.p for a in gen_accepted) + self.gen_accepted = list(gen_accepted) + values = [1] * self.num_contexts + list(gen_accepted) + self.accepted.copy_(torch.tensor(values, dtype=torch.int32)) + + +def _shared_planes(index0, position0): + """Live copies of the planes as views of one byte buffer per page: every layer's index-K + page, then the position page (what the QSA cache manager packs when an NVFP4 pool's data + and scale pages would otherwise absorb them).""" + index_bytes = TPB * INDEX_DIM * 2 + page = torch.zeros(PAGES, LAYERS * index_bytes + TPB * 3 * 4, dtype=torch.uint8, device=DEV) + index = [ + page[:, layer * index_bytes : (layer + 1) * index_bytes] + .view(torch.bfloat16) + .view(PAGES, TPB, 1, INDEX_DIM) + for layer in range(LAYERS) + ] + position = page[:, LAYERS * index_bytes :].view(torch.int32).view(PAGES, TPB, 3) + assert index[0].stride(0) == page.stride(0) // 2 and not index[0].is_contiguous() + for live, pre in zip(index, index0): + live.copy_(pre) + position.copy_(position0) + return index, position + + +def _qsl(lengths): + starts = [0] + for length in lengths: + starts.append(starts[-1] + length) + return torch.tensor(starts, dtype=torch.int64, device=DEV) + + +def _write_planes(b, layer_index, position, req_idx, logical, keys, anchor_keys, coords): + """The index store, then a compression stand-in that overwrites each anchor + row's own slot -- the two writers of an attention layer's planes.""" + qsa_paged_index_store(keys, coords, layer_index, position, req_idx, logical, b.block_table, TPB) + pages = b.block_table[req_idx.long(), logical // TPB].long() + within = logical % TPB + anchor = ((logical + 1) % COMPRESS == 0).reshape(-1, 1, 1) + layer_index[pages, within] = torch.where(anchor, anchor_keys, layer_index[pages, within]) + + +def _verify_step(b, rollback): + """The forward's write sequence over all `p` rows of every generation + request, staged at the forward's call sites, then the worker's commit.""" + rollback.discard() + qsl = b.qsl + rollback.begin( + num_contexts=b.num_contexts, + prefix_tokens=b.prefix, + num_gens=b.n, + rows_per_gen=b.p, + slots=b.slots[b.num_contexts :], + request_indices=b.req_idx[b.prefix :], + logical_positions=b.logical[b.prefix :], + block_table=b.block_table, + tokens_per_block=TPB, + ) + rollback.stage_ple_pools(b.ngram, b.conv) + window = ngram_token_history(b.ids, b.ngram, b.slots, b.carry, qsl, EOS, NGRAM) + rollback.stage_ple_rows(window, b.x) + dilated_short_conv_paged( + b.x.t(), b.weight, None, b.conv, qsl, b.slots, b.carry, DILATION, True, -1 + ) + for layer in range(LAYERS): + rollback.stage_index(b.index[layer], b.position) + _write_planes( + b, + b.index[layer], + b.position, + b.req_idx, + b.logical, + b.keys[layer], + b.anchor_keys[layer], + b.coords, + ) + rollback.commit_speculative_states(b.accepted, b.slots, b.num_contexts) + + +def _reference(b): + """The state after running the same ops from the pre-step state over the + context rows and each generation request's accepted prefix only.""" + keep = list(range(b.prefix)) + lengths = list(b.ctx_lens) + for s, a in enumerate(b.gen_accepted): + first = b.prefix + s * b.p + keep += range(first, first + a) + lengths.append(a) + keep = torch.tensor(keep, dtype=torch.long, device=DEV) + qsl = _qsl(lengths) + ngram, conv = b.ngram0.clone(), b.conv0.clone() + index = [t.clone() for t in b.index0] + position = b.position0.clone() + ngram_token_history(b.ids[keep], ngram, b.slots, b.carry, qsl, EOS, NGRAM) + dilated_short_conv_paged( + b.x[keep].t(), b.weight, None, conv, qsl, b.slots, b.carry, DILATION, True, -1 + ) + for layer in range(LAYERS): + _write_planes( + b, + index[layer], + position, + b.req_idx[keep], + b.logical[keep], + b.keys[layer][keep], + b.anchor_keys[layer][keep], + b.coords[keep], + ) + return ngram, conv, index, position + + +def _bits(t): + return t.view(torch.int16) if t.dtype == torch.bfloat16 else t + + +def _assert_state(b, expected, what): + ngram, conv, index, position = expected + assert torch.equal(b.ngram, ngram), f"{what}: n-gram pool" + assert torch.equal(_bits(b.conv), _bits(conv)), f"{what}: conv pool" + for layer in range(LAYERS): + assert torch.equal(_bits(b.index[layer]), _bits(index[layer])), f"{what}: index {layer}" + assert torch.equal(b.position, position), f"{what}: position plane" + + +# (ctx_lens, gen_starts, p, gen_accepted). Positions straddle pages and land on +# compression anchors; one request accepts everything. +CASES = { + "decode": ((), (5, 30, 61, 100), 4, (1, 2, 4, 3)), + "decode_p2": ((), (0, 31, 64), 2, (1, 2, 1)), + "mixed": ((7, 3), (12, 29, 95), 4, (2, 1, 3)), + "mixed_p3": ((40,), (2, 62), 3, (3, 1)), + # max_draft_len 7: every acceptance 1..8 once; blocks at 27, 60 and 90 straddle a page, + # and an 8-row block completes two compression groups. + "decode_p8": ((), (0, 27, 60, 90, 101, 5, 44, 119), 8, (1, 2, 3, 4, 5, 6, 7, 8)), + "mixed_p8": ((7, 3), (12, 29, 95), 8, (8, 1, 5)), +} + + +def _slots(num): + return list(range(3, 3 + 2 * num, 2)) + + +@pytest.mark.parametrize("case", sorted(CASES)) +def test_commit_matches_the_accepted_prefix(case): + ctx_lens, gen_starts, p, gen_accepted = CASES[case] + b = _Batch(11, ctx_lens, gen_starts, p, _slots(len(ctx_lens) + len(gen_starts))) + b.set_accepted(gen_accepted) + _verify_step(b, SpeculativeRollback()) + _assert_state(b, _reference(b), case) + + +@pytest.mark.parametrize("case", ["mixed", "decode_p8"]) +def test_commit_over_the_strided_planes_of_a_shared_page(case): + """The handler stages and restores the index and position planes by advanced indexing, so + planes that are strided views of one buffer (an NVFP4 KV cache's side planes) roll back + like contiguous ones, and the bytes of the other layers' pages in the same buffer stay as + the step wrote them.""" + ctx_lens, gen_starts, p, gen_accepted = CASES[case] + b = _Batch( + 13, ctx_lens, gen_starts, p, _slots(len(ctx_lens) + len(gen_starts)), shared_page=True + ) + b.set_accepted(gen_accepted) + _verify_step(b, SpeculativeRollback()) + _assert_state(b, _reference(b), case) + + +@pytest.mark.parametrize("accepted", ["all", "one"]) +def test_commit_at_the_acceptance_extremes(accepted): + ctx_lens, gen_starts, p, _ = CASES["mixed"] + b = _Batch(12, ctx_lens, gen_starts, p, _slots(len(ctx_lens) + len(gen_starts))) + b.set_accepted([p if accepted == "all" else 1] * len(gen_starts)) + _verify_step(b, SpeculativeRollback()) + _assert_state(b, _reference(b), accepted) + + +def test_commit_without_staging_is_a_no_op(): + ctx_lens, gen_starts, p, gen_accepted = CASES["decode"] + b = _Batch(14, ctx_lens, gen_starts, p, _slots(len(gen_starts))) + b.set_accepted(gen_accepted) + rollback = SpeculativeRollback() + rollback.commit_speculative_states(b.accepted, b.slots, 0) + _assert_state(b, (b.ngram0, b.conv0, b.index0, b.position0), "unstaged") + + +def test_discard_drops_a_stale_staging(): + """A staged step whose commit never ran is dropped by the next forward.""" + ctx_lens, gen_starts, p, gen_accepted = CASES["decode"] + b = _Batch(15, ctx_lens, gen_starts, p, _slots(len(gen_starts))) + b.set_accepted(gen_accepted) + rollback = SpeculativeRollback() + rollback.begin( + num_contexts=0, + prefix_tokens=0, + num_gens=b.n, + rows_per_gen=p, + slots=b.slots, + request_indices=b.req_idx, + logical_positions=b.logical, + block_table=b.block_table, + tokens_per_block=TPB, + ) + rollback.discard() + rollback.commit_speculative_states(b.accepted, b.slots, 0) + _assert_state(b, (b.ngram0, b.conv0, b.index0, b.position0), "discarded") + + +def test_commit_is_sync_free(): + ctx_lens, gen_starts, p, gen_accepted = CASES["mixed"] + b = _Batch(16, ctx_lens, gen_starts, p, _slots(len(ctx_lens) + len(gen_starts))) + b.set_accepted(gen_accepted) + rollback = SpeculativeRollback() + torch.cuda.synchronize() + torch.cuda.set_sync_debug_mode("error") + try: + _verify_step(b, rollback) + finally: + torch.cuda.set_sync_debug_mode(0) + _assert_state(b, _reference(b), "sync-free") + + +def test_padding_duplicates_leave_the_unique_request_exact(): + # Requests 2 and 3 are the same padding dummy as request 1: one slot, one + # block-table row, written concurrently with no defined winner, so only the + # unique request is compared. + gen_starts, p = (5, 40, 40, 40), 4 + slots = [3, 7, 7, 7] + b = _Batch(17, (), gen_starts, p, slots, dup_of={2: 1, 3: 1}) + b.set_accepted((2, 3, 3, 3)) + _verify_step(b, SpeculativeRollback()) + # The unique request's slot is exactly its accepted prefix. + pages = b.block_table[0].long() + expected = _reference(_single(b, 0)) + assert torch.equal(b.ngram[3], expected[0][3]) + assert torch.equal(_bits(b.conv[3]), _bits(expected[1][3])) + for layer in range(LAYERS): + assert torch.equal(_bits(b.index[layer][pages]), _bits(expected[2][layer][pages])), ( + f"index {layer}" + ) + assert torch.equal(b.position[pages], expected[3][pages]) + + +def _single(b, s): + """`b` restricted to its generation request `s`, same pre-step state.""" + one = _Batch.__new__(_Batch) + one.__dict__.update(b.__dict__) + rows = torch.arange(b.prefix + s * b.p, b.prefix + (s + 1) * b.p, device=DEV) + one.ctx_lens, one.num_contexts, one.n, one.prefix = [], 0, 1, 0 + one.lengths = [b.p] + one.slots = b.slots[b.num_contexts + s : b.num_contexts + s + 1] + one.carry = b.carry[b.num_contexts + s : b.num_contexts + s + 1] + one.qsl = _qsl(one.lengths) + one.ids, one.x = b.ids[rows], b.x[rows] + one.keys = [k[rows] for k in b.keys] + one.anchor_keys = [k[rows] for k in b.anchor_keys] + one.coords = b.coords[rows] + one.logical = b.logical[rows] + one.req_idx = torch.zeros_like(b.req_idx[rows]) + one.block_table = b.block_table[b.num_contexts + s : b.num_contexts + s + 1] + one.gen_accepted = [b.gen_accepted[s]] + return one + + +#: Per case, the acceptance counts one captured graph is replayed with (the accepted count +#: of every generation request, bonus token included). +_REPLAYS = { + "decode": [(1, 1, 1, 1), (2, 4, 1, 3), (4, 3, 2, 1), (3, 3, 3, 3)], + "decode_p8": [ + (1,) * 8, + (8, 7, 6, 5, 4, 3, 2, 1), + (2, 8, 1, 7, 3, 6, 4, 5), + (8,) * 8, + ], +} + + +@pytest.mark.parametrize("case", sorted(_REPLAYS)) +def test_capture_and_replay_follow_the_accepted_counts(case): + """The whole staged step and commit captured once, replayed with the + accepted counts and the pre-step state rewritten in place.""" + ctx_lens, gen_starts, p, _ = CASES[case] + b = _Batch(18, ctx_lens, gen_starts, p, _slots(len(gen_starts))) + rollback = SpeculativeRollback() + b.set_accepted((p,) * len(gen_starts)) + + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + _verify_step(b, rollback) + torch.cuda.current_stream().wait_stream(side) + b.reset() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _verify_step(b, rollback) + + for accepted in _REPLAYS[case]: + b.reset() + b.set_accepted(accepted) + graph.replay() + torch.cuda.synchronize() + _assert_state(b, _reference(b), f"replay {accepted}") diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_vision_tower_calls.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_vision_tower_calls.py new file mode 100644 index 000000000000..d1a339bd5e8a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_vision_tower_calls.py @@ -0,0 +1,194 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qwen4_exp vision tower's call grouping. + +A request's frames run in calls of whole frames under `_VISION_CALL_MAX_PATCHES` rows. +The tower here has the checkpoint's geometry (1152 wide, 16 heads of 72, 2x2 merge) with +two random blocks, so every op the real tower runs runs here; the cap is a module +constant, lowered so a few hundred rows span several calls. The reference is the same +tower in one call. + +Cutting between frames leaves attention, the position embedding, the rope and every +row-wise op unchanged; what can differ is which kernel a call picks (a GEMM's split, the +fused attention kernel for a call whose longest frame is at most 64 rows), within the ops' +reference bands. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 + modeling as M, +) + +assert torch.cuda.is_available(), "the vision tower requires a CUDA device" + +HIDDEN = 1152 +HEADS = 16 +HEAD_DIM = HIDDEN // HEADS +INTER = 4304 +MERGE = 2 +MERGE_DIM = HIDDEN * MERGE * MERGE +OUT = 2560 +POS_ROWS = 48 * 48 +PIXEL_WIDTH = 3 * 2 * 16 * 16 +DEPTH = 2 +DT = torch.bfloat16 + + +def _core(): + g = torch.Generator(device="cuda").manual_seed(0) + + def w(*shape, scale=0.02): + return (scale * torch.randn(*shape, generator=g, device="cuda")).to(DT) + + def norm_w(): + return (1.0 + 0.1 * torch.randn(HIDDEN, generator=g, device="cuda")).to(DT) + + weights = { + "embed": torch.zeros(1, dtype=DT, device="cuda"), + "vis_patch_w": w(HIDDEN, 3, 2, 16, 16), + "vis_patch_b": w(HIDDEN), + "vis_pos": w(POS_ROWS, HIDDEN, scale=0.5), + "vis_m_n_w": norm_w(), + "vis_m_n_b": w(HIDDEN), + "vis_m_fc1_w": w(MERGE_DIM, MERGE_DIM), + "vis_m_fc1_b": w(MERGE_DIM), + "vis_m_fc2_w": w(OUT, MERGE_DIM), + "vis_m_fc2_b": w(OUT), + } + for b in range(DEPTH): + weights.update( + { + f"l{b}_vis_n1_w": norm_w(), + f"l{b}_vis_n1_b": w(HIDDEN), + f"l{b}_vis_qkv_w": w(3 * HIDDEN, HIDDEN), + f"l{b}_vis_qkv_b": w(3 * HIDDEN), + f"l{b}_vis_proj_w": w(HIDDEN, HIDDEN), + f"l{b}_vis_proj_b": w(HIDDEN), + f"l{b}_vis_n2_w": norm_w(), + f"l{b}_vis_n2_b": w(HIDDEN), + f"l{b}_vis_fc1_w": w(INTER, HIDDEN), + f"l{b}_vis_fc1_b": w(INTER), + f"l{b}_vis_fc2_w": w(HIDDEN, INTER), + f"l{b}_vis_fc2_b": w(HIDDEN), + } + ) + return SimpleNamespace( + w=weights, + v_depth=DEPTH, + v_hidden=HIDDEN, + v_heads=HEADS, + v_head_dim=HEAD_DIM, + v_merge=MERGE, + v_merge_dim=MERGE_DIM, + v_pixel_width=PIXEL_WIDTH, + ) + + +def _pixels(grid, seed=1): + patches = sum(t * h * w for t, h, w in grid) + g = torch.Generator(device="cuda").manual_seed(seed) + return torch.empty(patches, PIXEL_WIDTH, device="cuda").uniform_(-1, 1, generator=g).to(DT) + + +def _count_calls(tower, monkeypatch): + """Wrap `tower._run`; the returned list gets one `(rows, grid)` entry per call.""" + calls = [] + run = tower._run + + def counted(pixel_values, grid, grid_thw): + calls.append((pixel_values.shape[0], list(grid))) + return run(pixel_values, grid, grid_thw) + + monkeypatch.setattr(tower, "_run", counted) + return calls + + +def _assert_close(out, ref): + """Two runs of the tower over different row counts differ by bf16 rounding: another + call picks another GEMM kernel, and the difference compounds through the blocks and the + merger. Measured here: 0.6% RMS relative error and at most 8 ulps (of the RMS) in one + element, against 117% RMS between unrelated rows. The gates sit between the two.""" + assert out.dtype == ref.dtype and out.shape == ref.shape, (out.shape, ref.shape) + error = out.float() - ref.float() + rms = ref.float().pow(2).mean().sqrt() + assert error.pow(2).mean().sqrt() <= 0.02 * rms, (error.pow(2).mean().sqrt(), rms) + assert error.abs().max() <= 16 * 2**-8 * rms, (error.abs().max(), rms) + + +@pytest.fixture(scope="module") +def tower(): + return M.VisionTower(_core()) + + +# `(grid, caps)`: frames of 256, 144 and 32 rows, of 400, 96 and 60, of 32 and 8, and of 384. +# Several caps per grid put the frame boundaries at different places, and the short-frame grid +# makes calls whose longest frame is at most 64 rows, where the fused attention kernel is the +# short-frame one. +CASES = { + "mixed": ([(2, 16, 16), (3, 12, 12), (6, 4, 8)], (512, 700)), + "image_then_video": ([(1, 20, 20), (4, 8, 12), (1, 6, 10)], (512, 700)), + "short_frames_only": ([(9, 4, 8), (3, 2, 4)], (96, 100)), + "one_item_many_frames": ([(7, 16, 24)], (800, 1200)), +} + + +@pytest.mark.parametrize( + "name,cap", [(name, cap) for name, (_, caps) in CASES.items() for cap in caps] +) +def test_calls_of_whole_frames_equal_one_call(tower, monkeypatch, name, cap): + grid = CASES[name][0] + pixels = _pixels(grid) + grid_thw = torch.tensor(grid) + monkeypatch.setattr(M, "_VISION_CALL_MAX_PATCHES", 1 << 20) + whole = tower(pixels, grid_thw) + + monkeypatch.setattr(M, "_VISION_CALL_MAX_PATCHES", cap) + calls = _count_calls(tower, monkeypatch) + chunked = tower(pixels, grid_thw) + + expected = M._frame_calls([tuple(row) for row in grid], cap) + assert len(calls) == len(expected) > 1 + assert all(rows <= cap for rows, _ in calls) + _assert_close(chunked, whole) + + +def test_a_request_that_fits_is_one_call_bit_for_bit(tower, monkeypatch): + grid = CASES["mixed"][0] + pixels = _pixels(grid) + direct = tower._run(pixels, [tuple(row) for row in grid], torch.tensor(grid)) + monkeypatch.setattr(M, "_VISION_CALL_MAX_PATCHES", 1136) + calls = _count_calls(tower, monkeypatch) + out = tower(pixels, torch.tensor(grid)) + assert [rows for rows, _ in calls] == [1136] + assert out.shape == (1136 // 4, OUT) + torch.testing.assert_close(out, direct, rtol=0, atol=0) + + +def test_a_frame_past_the_ceiling_is_refused_before_any_call(tower, monkeypatch): + monkeypatch.setattr(M, "_VISION_CALL_MAX_PATCHES", 256) + calls = _count_calls(tower, monkeypatch) + grid = [(1, 8, 8), (1, 20, 20)] + with pytest.raises(AssertionError, match="one frame is 20x20 = 400 patch rows"): + tower(_pixels(grid), torch.tensor(grid)) + assert calls == [] + + +def test_a_request_past_the_old_whole_request_ceiling_runs(tower): + # 66 frames of 512x512 pixels: 33 temporal groups of 32x32 patches, 33792 rows in two + # calls (32 groups, then 1). The last group's rows, alone in its call, and the first + # group's, in a call with 31 others, equal each group run on its own. + grid = [(33, 32, 32)] + pixels = _pixels(grid) + out = tower(pixels, torch.tensor(grid)) + assert out.shape == (33 * 32 * 32 // 4, OUT) + assert torch.isfinite(out).all() + group_rows, group_out = 32 * 32, 32 * 32 // 4 + for group in (0, 32): + alone = tower( + pixels[group * group_rows : (group + 1) * group_rows], torch.tensor([(1, 32, 32)]) + ) + _assert_close(out[group * group_out : (group + 1) * group_out], alone) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_kv_quant.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_kv_quant.py new file mode 100644 index 000000000000..6ca70bcdc1c2 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_kv_quant.py @@ -0,0 +1,245 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""qwen4_exp derives its KV storage from the engine's `kv_cache_quant_algo`. + +`None` keeps the checkpoint's bf16 K/V with the attention op's `quant_mode` 0 and no KV +scales; `FP8` stores K/V as e4m3 under `QM_FP8_KV_CACHE` with unit fp32 `[1]` scales; `NVFP4` +stores e2m1 pairs as uint8 half a head wide under `QM_NVFP4_KV_CACHE` with unit fp32 `[3]` +scales; any other algorithm is refused at construction. The core is built on the meta device from the +checkpoint's text and vision config values, so nothing here touches a device; the +`post_load_weights` cases stub out everything that reads weight values. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest import mock + +import pytest +import torch + +__extra_import_path__ = [".."] # noqa: F841 -- repo's file-scoped import hook + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 + modeling as M, +) +from tensorrt_llm._torch.configs import Qwen4ExpTextConfig, Qwen4ExpVisionConfig +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm.llmapi import QSASparseAttentionConfig +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +pytestmark = pytest.mark.cpu_only + +NUM_LAYERS = 48 + +# Qwen3.8-Flash-Next's text config. +_TEXT = dict( + architectures=["Qwen4ExpForCausalLM"], + attention_bias=False, + attention_dropout=0.0, + bos_token_id=248044, + dtype="bfloat16", + eos_token_id=248044, + full_attention_interval=4, + hc_count=4, + hc_lowrank=320, + head_dim=256, + heads_per_ngram=8, + hidden_act="silu", + hidden_size=2560, + indexer_budget=2048, + indexer_compress_ratio=4, + indexer_head_dim=128, + indexer_kv_heads=1, + indexer_n_heads=4, + initializer_range=0.02, + layer_types=["full_attention" if i % 4 == 3 else "linear_attention" for i in range(NUM_LAYERS)], + linear_conv_kernel_dim=4, + linear_key_head_dim=128, + linear_num_key_heads=16, + linear_num_value_heads=48, + linear_value_head_dim=128, + make_ngram_vocab_size_divisible_by=128, + mamba_ssm_dtype="float32", + max_position_embeddings=262144, + moe_intermediate_size=640, + mtp=dict( + hybrid=True, + layer_types=["full_attention"], + mtp_use_hidden_state_from_layer=None, + num_hidden_layers=1, + rope_theta=10000000, + ), + mtp_num_hidden_layers=1, + mtp_use_dedicated_embeddings=False, + ngram_size=3, + ngram_vocab_size_base=20000000, + norm_topk_prob=True, + num_attention_heads=24, + num_experts=512, + num_experts_per_tok=10, + num_hidden_layers=NUM_LAYERS, + num_key_value_heads=2, + number_of_conv_states=3, + output_gate_type="sigmoid", + output_router_logits=False, + pad_token_id=None, + partial_rotary_factor=0.25, + ple_conv_kernel_size=4, + ple_embed_dim=2560, + ple_layer_ids=[2], + rms_norm_eps=1e-06, + rope_parameters=dict( + mrope_interleaved=True, + mrope_section=[11, 11, 10], + partial_rotary_factor=0.25, + rope_theta=10000000, + rope_type="default", + ), + router_aux_loss_coef=0.001, + seed=1234, + shared_expert_intermediate_size=640, + split_ngram_parts=128, + tie_word_embeddings=False, + use_cache=True, + vocab_size=248320, +) +_VISION = dict( + deepstack_visual_indexes=[], + depth=27, + dtype="bfloat16", + hidden_act="gelu_pytorch_tanh", + hidden_size=1152, + in_channels=3, + intermediate_size=4304, + num_heads=16, + num_position_embeddings=2304, + out_hidden_size=2560, + patch_size=16, + spatial_merge_size=2, + temporal_patch_size=2, +) + + +def _core(kv_cache_quant_algo): + """The core of a Qwen3.8-Flash-Next engine whose resolved `kv_cache_quant_algo` is + `kv_cache_quant_algo`, with every parameter on the meta device.""" + model_config = ModelConfig( + pretrained_config=Qwen4ExpTextConfig(**_TEXT), + mapping=Mapping(world_size=1, tp_size=1, rank=0), + quant_config=QuantConfig( + quant_algo=QuantAlgo.MIXED_PRECISION, kv_cache_quant_algo=kv_cache_quant_algo + ), + moe_backend="CUTLASS", + moe_backend_requested="TRTLLM", + sparse_attention_config=QSASparseAttentionConfig(), + ) + composite = SimpleNamespace( + vision_config=Qwen4ExpVisionConfig(**_VISION), image_token_id=248056, video_token_id=248057 + ) + # The core asserts the device it runs on; the weights it declares stay on meta. + with ( + mock.patch.object(torch.cuda, "get_device_capability", return_value=(10, 3)), + torch.device("meta"), + ): + return M.Qwen4ExpTp1ModelingV2(model_config, composite) + + +def _load(core, monkeypatch) -> None: + """`post_load_weights` on a core with the weights it reads replaced by CPU stubs. + + Everything that reads weight values (the expert scalars, the n-gram table's driver + residency, the per-step targets) is replaced; what is left is the derivation under test + and the shape checks around it. The derived scales land on `embed`'s device, here CPU, + so their values can be read. + """ + table_rows = core.ngram_vocab_divisor + stubs = dict( + embed=torch.empty(1), + ple_table=torch.empty(table_rows, 1), + ple_offs=torch.zeros(1, dtype=torch.int64), + ple_vocab=torch.full((1,), table_rows, dtype=torch.int64), + ) + for key, value in stubs.items(): + core.w[key] = torch.nn.Parameter(value, requires_grad=False) + monkeypatch.setattr(M, "_moe_layer_scalars", lambda w, i, backend: {}) + for target in ("PrefillTarget", "DecodeTarget", "VisionTower"): + monkeypatch.setattr(M, target, lambda core: None) + monkeypatch.setattr(core, "_report_ple_residency", lambda: None) + core.post_load_weights() + + +def test_unquantized_kv_cache_keeps_the_checkpoints_bf16() -> None: + core = _core(None) + assert core.quant_mode == 0 + assert core.kv_pool_dtype == torch.bfloat16 + assert core._pool_geometry(32).kv_dtype == torch.bfloat16 + assert core.kv_scale_orig_quant is None and core.kv_scale_quant_orig is None + + +def test_fp8_kv_cache_stores_e4m3_under_the_fp8_quant_mode() -> None: + core = _core(QuantAlgo.FP8) + assert core.quant_mode == M.QM_FP8_KV_CACHE + assert core.kv_pool_dtype == torch.float8_e4m3fn + # The pool contract holds every manager's K/V pool to this dtype; the index and position + # planes keep the activation dtype. + geometry = core._pool_geometry(32) + assert geometry.kv_dtype == torch.float8_e4m3fn + assert geometry.dtype == torch.bfloat16 + + +def test_nvfp4_kv_cache_stores_e2m1_pairs_under_the_nvfp4_quant_mode() -> None: + core = _core(QuantAlgo.NVFP4) + assert core.quant_mode == M.QM_NVFP4_KV_CACHE + # An NVFP4 byte holds two values: the pool is uint8, half a head wide. + assert core.kv_pool_dtype == torch.uint8 + geometry = core._pool_geometry(32) + assert geometry.kv_dtype == torch.uint8 + assert geometry.kv_width == core.head_dim // 2 + assert geometry.dtype == torch.bfloat16 + + +@pytest.mark.parametrize("algo", [None, QuantAlgo.FP8]) +def test_unquantized_and_fp8_pools_are_a_full_head_wide(algo) -> None: + core = _core(algo) + assert core._pool_geometry(32).kv_width == core.head_dim + + +@pytest.mark.parametrize("algo", [QuantAlgo.INT8, QuantAlgo.FP8_BLOCK_SCALES]) +def test_other_kv_cache_algorithms_are_refused(algo) -> None: + """The QSA backend serves its sparse path over half, bf16, fp8 and NVFP4 pools and falls + back to dense for any other, so a pool of another storage must not be accepted.""" + with pytest.raises(AssertionError, match="kv_cache_quant_algo"): + _core(algo) + + +def test_fp8_kv_cache_scales_are_unit_fp32_single_elements(monkeypatch) -> None: + """The generation kernels read `kvCacheScale[0]` and the QSA store and sparse GQA take + no scale, so the scale is a real tensor holding exactly 1.0: the dense and the sparse + path then agree on what a cache slot holds.""" + core = _core(QuantAlgo.FP8) + _load(core, monkeypatch) + for scale in (core.kv_scale_orig_quant, core.kv_scale_quant_orig): + assert scale is not None + assert scale.dtype == torch.float32 and tuple(scale.shape) == (1,) + assert torch.equal(scale, torch.ones(1)) + + +def test_nvfp4_kv_cache_scales_are_unit_fp32_triples(monkeypatch) -> None: + """The attention op reads the q, K and V second-level scales out of one fp32 `[3]` tensor, + and the QSA NVFP4 store and sparse GQA take the same tensor, so every writer and reader + applies the one value, 1.0 (the checkpoint has no KV scales).""" + core = _core(QuantAlgo.NVFP4) + _load(core, monkeypatch) + for scale in (core.kv_scale_orig_quant, core.kv_scale_quant_orig): + assert scale is not None + assert scale.dtype == torch.float32 and tuple(scale.shape) == (3,) + assert torch.equal(scale, torch.ones(3)) + + +def test_bf16_kv_cache_builds_no_scales(monkeypatch) -> None: + core = _core(None) + _load(core, monkeypatch) + assert core.kv_scale_orig_quant is None and core.kv_scale_quant_orig is None diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_moe_scalars.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_moe_scalars.py new file mode 100644 index 000000000000..a58994c365a7 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_moe_scalars.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""qwen4_exp's routed-expert scalars, derived on CPU from the checkpoint's scale rows. + +CuteDSL requantizes FC1's output with ONE global scale per layer, where trtllm-gen folds +each expert's `down_proj.input_scale` into its own FC1 scalar. The CuteDSL scalars are +exact for any spread of `down_proj.input_scale` across experts: they are the engine's +`max(w2 input_scale)` convention. +""" + +from __future__ import annotations + +import pytest +import torch + +__extra_import_path__ = [".."] # noqa: F841 -- repo's file-scoped import hook + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 + modeling as M, +) + +EXPERTS = 8 + + +def _rows(isc2: torch.Tensor, seed: int = 0) -> dict: + gen = torch.Generator().manual_seed(seed) + isc1 = torch.full((EXPERTS,), 0.0123) + ws2_1 = torch.rand(EXPERTS, generator=gen) * 1e-3 + 1e-4 + return { + "l0_e_isc1": isc1, + "l0_e_isc1_up": isc1.clone(), + "l0_e_ws2_1": ws2_1, + "l0_e_ws2_1_up": ws2_1.clone(), + "l0_e_isc2": isc2, + "l0_e_ws2_2": torch.rand(EXPERTS, generator=gen) * 1e-3 + 1e-4, + } + + +def test_uniform_isc2_keeps_the_per_expert_products_bitwise() -> None: + w = _rows(torch.full((EXPERTS,), 0.0371)) + got = M._moe_layer_scalars(w, 0, "CUTEDSL") + assert torch.equal(got["alpha2"], w["l0_e_isc2"] * w["l0_e_ws2_2"]) + assert torch.equal(got["global_sf"], 1.0 / w["l0_e_isc2"][:1]) + assert got["global_sf"].shape == (1,) + assert torch.equal(got["alpha1"], w["l0_e_isc1"] * w["l0_e_ws2_1"]) + + +@pytest.mark.parametrize("spread", [1.0, 1.0001, 3.0, 40.0]) +def test_isc2_spread_is_folded_into_alpha2(spread: float) -> None: + """Whatever the spread, `global_sf * alpha2[e]` is expert e's own `ws2_2[e]`, the + product trtllm-gen forms from `(1 / isc2[e]) * (isc2[e] * ws2_2[e])`.""" + isc2 = 0.01 * torch.linspace(1.0, spread, EXPERTS) + w = _rows(isc2.flip(0)) + cute = M._moe_layer_scalars(w, 0, "CUTEDSL") + gen = M._moe_layer_scalars(w, 0, "TRTLLM") + isc2 = w["l0_e_isc2"] + assert torch.equal(cute["global_sf"], 1.0 / isc2.max().reshape(1)) + # The CuteDSL FC1 quantizes at the smallest per-expert global scale: no expert's FC2 + # input is scaled above what the checkpoint calibrated for it. + assert (cute["global_sf"] <= 1.0 / isc2).all() + assert torch.allclose( + cute["global_sf"] * cute["alpha2"], (1.0 / isc2) * gen["output2_scale_scalar"], rtol=1e-6 + ) + assert torch.allclose(cute["global_sf"] * cute["alpha2"], w["l0_e_ws2_2"], rtol=1e-6) + + +def test_cutedsl_fc2_matches_the_per_expert_reference() -> None: + """The FC2 value an expert produces is `ws2_2[e] * (y @ W)` whether its input was + quantized at its own `1 / isc2[e]` (trtllm-gen) or at the layer's single scale.""" + isc2 = torch.tensor([0.02, 0.01, 0.04, 0.03, 0.02, 0.05, 0.01, 0.025]) + w = _rows(isc2) + cute = M._moe_layer_scalars(w, 0, "CUTEDSL") + gen = M._moe_layer_scalars(w, 0, "TRTLLM") + torch.manual_seed(0) + y = torch.randn(EXPERTS, 64, dtype=torch.float64) + weight = torch.randn(EXPERTS, 64, dtype=torch.float64) + ref = w["l0_e_ws2_2"].double() * (y * weight).sum(-1) + per_expert = gen["output2_scale_scalar"].double() * ( + (1.0 / isc2.double())[:, None] * y * weight + ).sum(-1) + one_scale = cute["alpha2"].double() * (cute["global_sf"].double() * y * weight).sum(-1) + torch.testing.assert_close(per_expert, ref, rtol=1e-5, atol=0) + torch.testing.assert_close(one_scale, ref, rtol=1e-5, atol=0) + + +def test_gate_up_input_scale_mismatch_is_refused() -> None: + w = _rows(torch.full((EXPERTS,), 0.0371)) + w["l0_e_isc1_up"] = w["l0_e_isc1_up"] * 1.5 + with pytest.raises(AssertionError, match="gate/up input_scale differ"): + M._moe_layer_scalars(w, 0, "CUTEDSL") diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_multimodal_embeds.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_multimodal_embeds.py new file mode 100644 index 000000000000..48d869234d32 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_multimodal_embeds.py @@ -0,0 +1,294 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""qwen4_exp's multimodal splice runs the vision tower once per modality a request +carries and puts the rows back in prompt order; the tower itself runs a request's frames +in calls of whole frames under a row ceiling. + +The tower is a stub whose rows name the item they came from, so the bookkeeping -- +modality dispatch, `mm_item_order`, the chunk slice, the frame grouping and the order of +the concatenated calls -- runs on the CPU without the vision ops or an engine. +""" + +from __future__ import annotations + +import types +from types import SimpleNamespace + +import pytest +import torch + +__extra_import_path__ = [".."] # noqa: F841 -- repo's file-scoped import hook + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 + modeling as M, +) + +pytestmark = pytest.mark.cpu_only + +Shell = M.ModelingV2Qwen38FlashNextNvfp4Sm103Tp1 + +HIDDEN = 3 +MERGE = 2 +IMAGE, VIDEO = 1, 2 + + +class _Tower: + """Stands in for `VisionTower`: every row of item `i` of modality `m` holds + `100 * m + i`, so a spliced row says where it came from.""" + + def __init__(self): + self.calls = 0 + + def __call__(self, pixel_values, grid_thw): + self.calls += 1 + modality = int(pixel_values[0, 0]) + return torch.cat( + [ + torch.full((t * h * w // (MERGE * MERGE), HIDDEN), float(100 * modality + item)) + for item, (t, h, w) in enumerate(grid_thw.tolist()) + ] + ) + + +def _shell(tower): + shell = SimpleNamespace( + model=SimpleNamespace( + v_merge=MERGE, w={"embed": torch.zeros(16, HIDDEN)}, vision_tower=tower + ) + ) + shell._run_tower = types.MethodType(Shell._run_tower, shell) + shell._in_prompt_order = types.MethodType(Shell._in_prompt_order, shell) + return shell + + +def _media(modality, grids): + """One `multimodal_data` entry; `pixel_values` only carries the stub's modality.""" + patches = sum(t * h * w for t, h, w in grids) + key = "pixel_values" if modality == IMAGE else "pixel_values_videos" + grid_key = "image_grid_thw" if modality == IMAGE else "video_grid_thw" + return {key: torch.full((patches, 1), float(modality)), grid_key: torch.tensor(grids)} + + +def _params(image=None, video=None, order=None, runtime=None, mrope=True): + data = {} + if mrope: + data["mrope_config"] = {"mrope_position_ids": torch.zeros(3, 1, 8, dtype=torch.int32)} + if image is not None: + data["image"] = _media(IMAGE, image) + if video is not None: + data["video"] = _media(VIDEO, video) + return SimpleNamespace(multimodal_data=data, mm_item_order=order, multimodal_runtime=runtime) + + +def _splice(tower, params, num_rows, padding=0): + """The tower's rows spliced, as the prefill target's break does, over a prompt whose + tokens 1..num_rows are the placeholders and which `padding` rows follow; returns the + rows that landed there, one value per row, or None when the step carries no media.""" + kwargs = { + "multimodal_params": [params], + "mm_token_indices": torch.arange(1, 1 + num_rows, dtype=torch.int64), + } + vision = Shell._multimodal_embeds(_shell(tower), kwargs) + if vision is None: + return None + total = num_rows + 2 + h = torch.full((total + padding, HIDDEN), -1.0) + target = SimpleNamespace(_step=SimpleNamespace(vision=vision, num_tokens=total)) + M.PrefillTarget._splice(target, h) + assert torch.count_nonzero(h[total:]) == 0, "padding rows must be zeroed" + assert torch.all(h[0] == -1.0) and torch.all(h[total - 1] == -1.0) + return h[1 : 1 + num_rows, 0].tolist() + + +def _item(modality, index): + return {"modality": modality, "index": index} + + +# Two videos of 4 and 8 rows, two images of 4 and 1. +VIDEOS = [(2, 2, 4), (2, 4, 4)] +IMAGES = [(1, 4, 4), (1, 2, 2)] +MIXED_ORDER = [_item("video", 0), _item("image", 0), _item("video", 1), _item("image", 1)] +MIXED_ROWS = [200.0] * 4 + [100.0] * 4 + [201.0] * 8 + [101.0] + + +def test_image_request_is_unchanged(): + tower = _Tower() + assert _splice(tower, _params(image=IMAGES), 5) == [100.0] * 4 + [101.0] + assert tower.calls == 1 + + +def test_the_splice_zeroes_the_padding_rows(): + assert _splice(_Tower(), _params(image=IMAGES), 5, padding=3) == [100.0] * 4 + [101.0] + + +def test_a_text_step_writes_only_the_padding(): + target = SimpleNamespace(_step=SimpleNamespace(vision=None, num_tokens=2)) + h = torch.full((4, HIDDEN), -1.0) + M.PrefillTarget._splice(target, h) + assert torch.all(h[:2] == -1.0) and torch.count_nonzero(h[2:]) == 0 + + +def test_video_rows_reach_the_placeholders_item_by_item(): + tower = _Tower() + assert _splice(tower, _params(video=VIDEOS), 12) == [200.0] * 4 + [201.0] * 8 + assert tower.calls == 1 + + +def test_a_mixed_request_is_reassembled_in_prompt_order(): + tower = _Tower() + params = _params(image=IMAGES, video=VIDEOS, order=MIXED_ORDER) + assert _splice(tower, params, len(MIXED_ROWS)) == MIXED_ROWS + assert tower.calls == 2 + + +def test_a_mixed_request_needs_its_manifest(): + with pytest.raises(AssertionError, match="needs mm_item_order"): + _splice(_Tower(), _params(image=IMAGES, video=VIDEOS), len(MIXED_ROWS)) + + +def test_a_manifest_that_disagrees_with_the_payload_is_refused(): + params = _params(image=IMAGES, video=VIDEOS, order=MIXED_ORDER[:3]) + with pytest.raises(AssertionError, match="mm_item_order lists image items"): + _splice(_Tower(), params, len(MIXED_ROWS)) + + +def test_a_chunk_takes_its_slice_of_the_prompt_order_rows(): + runtime = SimpleNamespace( + num_cached_mm_tokens=3, num_mm_tokens_in_chunk=6, total_embeds_in_request=len(MIXED_ROWS) + ) + params = _params(image=IMAGES, video=VIDEOS, order=MIXED_ORDER, runtime=runtime) + assert _splice(_Tower(), params, 6) == MIXED_ROWS[3:9] + + +def test_a_chunk_without_placeholders_skips_the_tower(): + runtime = SimpleNamespace( + num_cached_mm_tokens=len(MIXED_ROWS), + num_mm_tokens_in_chunk=0, + total_embeds_in_request=len(MIXED_ROWS), + ) + tower = _Tower() + params = _params(image=IMAGES, video=VIDEOS, order=MIXED_ORDER, runtime=runtime) + assert _splice(tower, params, 0) is None + assert tower.calls == 0 + + +def _frames(grid): + """The item-by-item frame list a grid describes: one `(h, w)` per frame.""" + return [(h, w) for t, h, w in grid for _ in range(t)] + + +def _random_grid(gen): + return [ + (int(torch.randint(1, 6, (1,), generator=gen)), 2 * h, 2 * w) + for h, w in torch.randint( + 1, 5, (int(torch.randint(1, 6, (1,), generator=gen)), 2), generator=gen + ).tolist() + ] + + +def test_frames_are_grouped_whole_and_in_order(): + # 16-row frames, 32 rows to a call: one item's frames split across calls. + assert M._frame_calls([(5, 4, 4)], 32) == [[(2, 4, 4)], [(2, 4, 4)], [(1, 4, 4)]] + # A call takes the end of one item and the start of the next. + grid = [(1, 8, 4), (3, 4, 4), (1, 2, 2)] + assert M._frame_calls(grid, 48) == [[(1, 8, 4), (1, 4, 4)], [(2, 4, 4), (1, 2, 2)]] + assert M._frame_calls(grid, 1 << 20) == [grid] + + +def test_a_call_is_filled_to_the_ceiling_and_never_past_it(): + gen = torch.Generator().manual_seed(0) + for _ in range(200): + grid = _random_grid(gen) + cap = int(torch.randint(64, 300, (1,), generator=gen)) + calls = M._frame_calls(grid, cap) + assert [f for call in calls for f in _frames(call)] == _frames(grid) + sizes = [sum(t * h * w for t, h, w in call) for call in calls] + assert all(0 < size <= cap for size in sizes) + # Greedy: the next call's first frame would not have fit. + for size, call in zip(sizes, calls[1:]): + h, w = _frames(call)[0] + assert size + h * w > cap + + +def test_one_frame_past_the_ceiling_is_refused_and_one_at_it_is_not(): + assert M._frame_calls([(2, 8, 8)], 64) == [[(1, 8, 8)], [(1, 8, 8)]] + with pytest.raises(AssertionError, match="one frame is 8x10 = 80 patch rows"): + M._frame_calls([(1, 4, 4), (3, 8, 10)], 64) + + +PIXEL_WIDTH = 4 + + +def _tower(monkeypatch, cap): + """A `VisionTower` whose per-call body is a stub: merged row `j` of a call carries the + pixel value of the call's patch row `4 * j`, which `_pixels` set to the patch's + index in the request.""" + monkeypatch.setattr(M, "_VISION_CALL_MAX_PATCHES", cap) + tower = object.__new__(M.VisionTower) + tower.core = SimpleNamespace( + w={"embed": torch.zeros(1)}, v_pixel_width=PIXEL_WIDTH, v_merge=MERGE + ) + tower.calls = [] + + def run(pixel_values, grid, grid_thw): + tower.calls.append((pixel_values, grid, grid_thw)) + return pixel_values[:: MERGE * MERGE, :1].expand(-1, HIDDEN).clone() + + tower._run = run + return tower + + +def _pixels(grid): + patches = sum(t * h * w for t, h, w in grid) + return torch.arange(patches, dtype=torch.float32)[:, None].expand(-1, PIXEL_WIDTH).clone() + + +def test_the_tower_calls_whole_frames_and_concatenates_their_rows_in_order(monkeypatch): + grid = [(1, 8, 4), (3, 4, 4), (2, 2, 6), (1, 2, 2)] + tower = _tower(monkeypatch, 48) + out = tower(_pixels(grid), torch.tensor(grid)) + patches = sum(t * h * w for t, h, w in grid) + assert out[:, 0].tolist() == [float(4 * j) for j in range(patches // 4)] + assert [call[1] for call in tower.calls] == M._frame_calls(grid, 48) + start = 0 + for pixel_values, call, grid_thw in tower.calls: + rows = sum(t * h * w for t, h, w in call) + assert pixel_values[:, 0].tolist() == [float(i) for i in range(start, start + rows)] + assert rows <= 48 + assert grid_thw.tolist() == [list(row) for row in call] + assert grid_thw.dtype == torch.int64 + start += rows + assert start == patches + + +def test_a_request_under_the_ceiling_is_one_call(monkeypatch): + grid = [(2, 4, 4), (1, 4, 8)] + tower = _tower(monkeypatch, 64) + tower(_pixels(grid), torch.tensor(grid)) + assert len(tower.calls) == 1 + assert tower.calls[0][2].tolist() == [list(row) for row in grid] + + +def test_a_request_past_the_old_whole_request_ceiling_is_served(monkeypatch): + # 33 temporal groups of 32x32 patches (66 frames of 512x512 pixels): 33792 rows, 1024 + # rows past what a single call may carry. + ceiling = M._VISION_CALL_MAX_PATCHES + grid = [(33, 32, 32)] + tower = _tower(monkeypatch, ceiling) + out = tower(_pixels(grid), torch.tensor(grid)) + assert [call[1] for call in tower.calls] == [[(32, 32, 32)], [(1, 32, 32)]] + assert out.shape[0] == 33 * 32 * 32 // 4 + assert out[:, 0].tolist() == [float(4 * j) for j in range(out.shape[0])] + + +def test_a_frame_past_the_ceiling_cannot_be_served(monkeypatch): + tower = _tower(monkeypatch, 64) + grid = [(1, 4, 4), (1, 10, 10)] + with pytest.raises(AssertionError, match="one frame is 10x10 = 100 patch rows"): + tower(_pixels(grid), torch.tensor(grid)) + assert tower.calls == [] + + +def test_a_media_request_needs_the_processors_position_ids(): + with pytest.raises(AssertionError, match="mrope_position_ids"): + _splice(_Tower(), _params(video=VIDEOS, mrope=False), 12) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_pool_contract.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_pool_contract.py index d89a528d1417..93e0bd76cc29 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_pool_contract.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_pool_contract.py @@ -9,15 +9,20 @@ from __future__ import annotations +from types import SimpleNamespace + import pytest import torch __extra_import_path__ = [".."] # noqa: F841 -- repo's file-scoped import hook +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp import routing from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 modeling as M, ) +pytestmark = pytest.mark.cpu_only + G = M._PoolGeometry( tokens_per_block=32, heads_kv=2, @@ -26,6 +31,7 @@ index_head_dim=128, conv_dim=8192, conv_kernel=4, + num_k_heads=16, num_v_heads=48, head_v_dim=128, head_k_dim=128, @@ -33,8 +39,15 @@ ple_state_len=9, ngram_context_len=3, dtype=torch.bfloat16, + kv_dtype=torch.bfloat16, + kv_width=256, ) +# The same pools under an fp8 KV cache: K/V are e4m3, every other plane keeps its type. +G_FP8 = G._replace(kv_dtype=torch.float8_e4m3fn) +# Under NVFP4 the K/V pool is uint8 e2m1 pairs, half a head wide. +G_NVFP4 = G._replace(kv_dtype=torch.uint8, kv_width=128) SLOTS = 6 +HISTORY = 16 def _t(*shape, dtype=torch.bfloat16): @@ -45,8 +58,17 @@ def _conv(**kw): return _t(SLOTS, G.conv_dim, G.conv_kernel - 1, **kw) -def _ssm(**kw): - return _t(SLOTS, G.num_v_heads, G.head_v_dim, G.head_k_dim, dtype=torch.float32) +def _ssm(dtype=torch.float32): + return _t(SLOTS, G.num_v_heads, G.head_v_dim, G.head_k_dim, dtype=dtype) + + +def _history(slots=SLOTS, history=HISTORY): + """One GDN layer's cached-replay history buffers (old_x, old_B, old_dt).""" + return ( + _t(slots, 2, history, G.num_v_heads, G.head_v_dim), + _t(slots, 2, history, G.num_k_heads, G.head_k_dim), + _t(slots, 2, G.num_v_heads, history, dtype=torch.float32), + ) def test_valid_pools_pass() -> None: @@ -54,7 +76,17 @@ def test_valid_pools_pass() -> None: M.check_attention_pools( G, 3, _t(9, 32, 1, 128), _t(9, 2, G.heads_kv, G.tokens_per_block, G.head_dim) ) - M.check_recurrent_pools(G, _conv(), _ssm()) + M.check_attention_pools( + G_FP8, + 3, + _t(9, 32, 1, 128), + _t(9, 2, G.heads_kv, G.tokens_per_block, G.head_dim, dtype=torch.float8_e4m3fn), + ) + nvfp4_kv = _t(9, 2, G.heads_kv, G.tokens_per_block, 128, dtype=torch.uint8) + M.check_attention_pools(G_NVFP4, 3, _t(9, 32, 1, 128), nvfp4_kv) + M.check_kv_scale_planes(G_NVFP4, 3, nvfp4_kv, _t(9, 2, 2, 32, 16, dtype=torch.uint8)) + for dtype in (torch.float32, torch.bfloat16): + M.check_recurrent_pools(G, _conv(), _ssm(dtype)) M.check_ple_pools( G, _t(SLOTS, G.stream, G.ple_state_len), @@ -62,13 +94,25 @@ def test_valid_pools_pass() -> None: _conv(), ) M.check_step_buffers( - G, - _t(SLOTS, 2, G.conv_dim, G.conv_kernel - 1), - _t(SLOTS, 2, G.num_v_heads, G.head_v_dim, G.head_k_dim, dtype=torch.float32), - _conv(), - _ssm(), - _t(SLOTS, dtype=torch.int32), + G, _t(SLOTS, 2, G.conv_dim, G.conv_kernel - 1), _conv(), _t(SLOTS, dtype=torch.int32) ) + for dtype in (torch.float32, torch.bfloat16): + M.check_legacy_state( + G, _t(SLOTS, 2, G.num_v_heads, G.head_v_dim, G.head_k_dim, dtype=dtype), _ssm(dtype) + ) + for dtype in (torch.float32, torch.bfloat16): + for width in (2, 5, M._MAX_VERIFY_WIDTH): + M.check_replay_state( + HISTORY, + width, + _t(SLOTS, dtype=torch.int32), + _t(SLOTS, dtype=torch.int32), + _ssm(dtype), + ) + M.check_replay_history(G, 3, HISTORY, SLOTS, *_history()) + for rows in range(1, M._MAX_VERIFY_WIDTH + 1): + M.check_verify_width(rows) + M.check_shared_selection(_t(2, 8, 512, dtype=torch.int32), {3: 1}, 3, 512) @pytest.mark.parametrize( @@ -101,6 +145,77 @@ def test_attention_pools_reject_wrong_index_plane_and_kv_layout() -> None: M.check_attention_pools(G, 3, _t(9, 32, 1, 128), _t(9, 3, 2, 32, 256)) +def test_attention_pools_reject_the_other_storage_dtype() -> None: + """The attention op reads the pool through one `quant_mode`, so a pool of the other + storage would be read through the wrong format. The draft layer's bind runs the same check.""" + index = _t(9, 32, 1, 128) + bf16_kv = _t(9, 2, G.heads_kv, G.tokens_per_block, G.head_dim) + fp8_kv = bf16_kv.to(torch.float8_e4m3fn) + with pytest.raises(AssertionError, match="paged K/V"): + M.check_attention_pools(G, 3, index, fp8_kv) + with pytest.raises(AssertionError, match="paged K/V"): + M.check_attention_pools(G, 3, index, bf16_kv.to(torch.float16)) + with pytest.raises(AssertionError, match="paged K/V"): + M.check_attention_pools(G_FP8, 3, index, bf16_kv) + # The index plane is a kernel contract, not KV storage: bf16 under either algorithm. + with pytest.raises(AssertionError, match="QSA index plane"): + M.check_attention_pools(G_FP8, 3, index.to(torch.float8_e4m3fn), fp8_kv) + + +def test_attention_pools_reject_the_wrong_nvfp4_width() -> None: + """An NVFP4 pool is half a head wide: a full-width uint8 pool is the other storage.""" + index = _t(9, 32, 1, 128) + with pytest.raises(AssertionError, match="paged K/V"): + M.check_attention_pools( + G_NVFP4, 3, index, _t(9, 2, G.heads_kv, G.tokens_per_block, 256, dtype=torch.uint8) + ) + with pytest.raises(AssertionError, match="paged K/V"): + M.check_attention_pools( + G, 3, index, _t(9, 2, G.heads_kv, G.tokens_per_block, 128, dtype=torch.uint8) + ) + + +def _strided(shape, slot_stride, dtype=torch.uint8): + """A view of `shape` whose slot stride is `slot_stride`, as a manager's role planes of a + lifecycle slot holding several layers' pages are.""" + pages = shape[0] * slot_stride + flat = torch.empty(pages, dtype=dtype, device="meta") + return flat.as_strided(shape, (slot_stride, *torch.empty(shape[1:], device="meta").stride())) + + +def test_kv_scale_planes_must_hold_the_data_pools_pages_per_slot() -> None: + """The attention op addresses the scale pages with the data pages' block table, so a + scale pool with another number of pages per slot would be read at other slots' scales. + The engine's trunk holds 24 data role pages per slot (12 layers x K, V): 196608 bytes + against 24576 of scales.""" + data = _strided((9, 2, G.heads_kv, G.tokens_per_block, 128), 24 * 8192) + scales = _strided((9, 2, G.heads_kv, G.tokens_per_block, 16), 24 * 1024) + M.check_kv_scale_planes(G_NVFP4, 3, data, scales) + # A side page coalesced into the data pool: 36 pages per slot against 24. + coalesced = _strided((9, 2, G.heads_kv, G.tokens_per_block, 128), 36 * 8192) + with pytest.raises(AssertionError, match="36 pages per slot.*24"): + M.check_kv_scale_planes(G_NVFP4, 3, coalesced, scales) + with pytest.raises(AssertionError, match="whole number"): + M.check_kv_scale_planes(G_NVFP4, 3, data, _strided((9, 2, 2, 32, 16), 24 * 1024 + 1)) + + +@pytest.mark.parametrize( + "scales, match", + [ + (None, "did not allocate"), + (_t(9, 2, 2, 32, 16), "torch.uint8"), + (_t(9, 2, 2, 32, 8, dtype=torch.uint8), "NVFP4 block scales"), + (_t(9, 2, 4, 32, 16, dtype=torch.uint8), "NVFP4 block scales"), + (_t(9, 3, 2, 32, 16, dtype=torch.uint8), "same slots"), + (_t(8, 2, 2, 32, 16, dtype=torch.uint8), "same slots"), + ], +) +def test_kv_scale_planes_reject(scales, match) -> None: + data = _t(9, 2, G.heads_kv, G.tokens_per_block, 128, dtype=torch.uint8) + with pytest.raises(AssertionError, match=match): + M.check_kv_scale_planes(G_NVFP4, 3, data, scales) + + def test_recurrent_pools_reject_wrong_key_dim_order_and_dtype() -> None: with pytest.raises(AssertionError, match="recurrent state pool"): M.check_recurrent_pools( @@ -108,6 +223,9 @@ def test_recurrent_pools_reject_wrong_key_dim_order_and_dtype() -> None: ) with pytest.raises(AssertionError, match="conv state pool"): M.check_recurrent_pools(G, _conv(dtype=torch.float32), _ssm()) + # The engine accepts an fp16 pool; no recurrence entry the target binds is certified at it. + with pytest.raises(AssertionError, match="recurrent state pool"): + M.check_recurrent_pools(G, _conv(), _ssm(torch.float16)) def test_ple_pools_reject_slot_mismatch_and_int32_history() -> None: @@ -123,19 +241,123 @@ def test_ple_pools_reject_slot_mismatch_and_int32_history() -> None: M.check_ple_pools(G, None, ngram, _conv()) -def test_step_buffers_reject_missing_dtype_drift_and_index() -> None: +def test_step_buffers_reject_missing_conv_rank_dtype_and_index() -> None: int_conv = _t(SLOTS, 2, G.conv_dim, G.conv_kernel - 1) - int_ssm = _t(SLOTS, 2, G.num_v_heads, G.head_v_dim, G.head_k_dim, dtype=torch.float32) idx = _t(SLOTS, dtype=torch.int32) - with pytest.raises(AssertionError, match="no per-step"): - M.check_step_buffers(G, None, None, _conv(), _ssm(), idx) + with pytest.raises(AssertionError, match="no per-step conv"): + M.check_step_buffers(G, None, _conv(), idx) # Rank without the steps dim. with pytest.raises(AssertionError, match="per-step conv buffer"): - M.check_step_buffers(G, _conv(), int_ssm, _conv(), _ssm(), idx) - # The promotion kernel copies without converting, so the dtypes must match the pools. - with pytest.raises(AssertionError, match="per-step recurrent buffer"): - M.check_step_buffers(G, int_conv, int_ssm.to(torch.bfloat16), _conv(), _ssm(), idx) + M.check_step_buffers(G, _conv(), _conv(), idx) + # The promotion kernel copies without converting, so the dtype must match the pool. + with pytest.raises(AssertionError, match="per-step conv buffer"): + M.check_step_buffers(G, int_conv.to(torch.float32), _conv(), idx) with pytest.raises(AssertionError, match="intermediate_state_indices"): - M.check_step_buffers(G, int_conv, int_ssm, _conv(), _ssm(), idx.to(torch.int64)) + M.check_step_buffers(G, int_conv, _conv(), idx.to(torch.int64)) with pytest.raises(AssertionError, match="intermediate_state_indices"): - M.check_step_buffers(G, int_conv, int_ssm, _conv(), _ssm(), None) + M.check_step_buffers(G, int_conv, _conv(), None) + + +def test_legacy_state_rejects_missing_rank_and_dtype_drift() -> None: + int_ssm = _t(SLOTS, 2, G.num_v_heads, G.head_v_dim, G.head_k_dim, dtype=torch.float32) + with pytest.raises(AssertionError, match="no per-step recurrent"): + M.check_legacy_state(G, None, _ssm()) + with pytest.raises(AssertionError, match="per-step recurrent buffer"): + M.check_legacy_state(G, _ssm(), _ssm()) + with pytest.raises(AssertionError, match="per-step recurrent buffer"): + M.check_legacy_state(G, int_ssm.to(torch.bfloat16), _ssm()) + + +def test_replay_state_rejects_pool_dtype_geometry_and_bookkeeping() -> None: + pnat = _t(SLOTS, dtype=torch.int32) + buf = _t(SLOTS, dtype=torch.int32) + with pytest.raises(AssertionError, match="fp32 or a bf16 checkpoint"): + M.check_replay_state(HISTORY, 3, pnat, buf, _ssm(torch.float16)) + with pytest.raises(AssertionError, match="replay_history_size"): + M.check_replay_state(12, 3, pnat, buf, _ssm()) + for width in (1, M._MAX_VERIFY_WIDTH + 1): + with pytest.raises(AssertionError, match="replay_step_width"): + M.check_replay_state(HISTORY, width, pnat, buf, _ssm()) + # The bookkeeping is one int32 per pool slot, so it must cover the whole pool. + with pytest.raises(AssertionError, match="prev_num_accepted_tokens"): + M.check_replay_state(HISTORY, 3, _t(SLOTS - 1, dtype=torch.int32), buf, _ssm()) + with pytest.raises(AssertionError, match="cache_buf_idx"): + M.check_replay_state(HISTORY, 3, pnat, buf.to(torch.int64), _ssm()) + + +def test_legacy_state_rejects_more_than_the_certified_rows_per_request() -> None: + def int_ssm(steps): + return _t(SLOTS, steps, G.num_v_heads, G.head_v_dim, G.head_k_dim, dtype=torch.float32) + + M.check_legacy_state(G, int_ssm(M._MAX_LEGACY_VERIFY_WIDTH), _ssm()) + with pytest.raises(AssertionError, match="TRTLLM_USE_GDN_REPLAY"): + M.check_legacy_state(G, int_ssm(M._MAX_LEGACY_VERIFY_WIDTH + 1), _ssm()) + + +def test_verify_width_admits_one_to_eight_rows_per_request() -> None: + for rows in (0, M._MAX_VERIFY_WIDTH + 1): + with pytest.raises(AssertionError, match="predicted_tokens_per_seq"): + M.check_verify_width(rows) + + +def test_replay_width_admits_only_the_managers_step_width() -> None: + replay = SimpleNamespace(width=3) + M.check_replay_width(None, 1) # the legacy recurrence takes any width + M.check_replay_width(replay, 3) + for rows in (1, 2, 4): + with pytest.raises(AssertionError, match="step width 3"): + M.check_replay_width(replay, rows) + + +def test_replay_history_rejects_missing_short_pool_and_wrong_geometry() -> None: + old_x, old_B, old_dt = _history() + with pytest.raises(AssertionError, match="no old_x"): + M.check_replay_history(G, 3, HISTORY, SLOTS, None, old_B, old_dt) + with pytest.raises(AssertionError, match="slots for a"): + M.check_replay_history(G, 3, HISTORY, SLOTS, *_history(slots=SLOTS - 1)) + # The kernel addresses the history by `history_size`, so a buffer of another depth + # would be read at the wrong offsets. + with pytest.raises(AssertionError, match="old_x"): + M.check_replay_history(G, 3, HISTORY, SLOTS, *_history(history=12)) + with pytest.raises(AssertionError, match="old_B"): + M.check_replay_history(G, 3, HISTORY, SLOTS, old_x, old_B.to(torch.float32), old_dt) + with pytest.raises(AssertionError, match="old_dt"): + M.check_replay_history(G, 3, HISTORY, SLOTS, old_x, old_B, old_dt.to(torch.bfloat16)) + with pytest.raises(AssertionError, match="double-buffered"): + M.check_replay_history( + G, + 3, + HISTORY, + SLOTS, + old_x, + old_B, + _t(SLOTS, 3, G.num_v_heads, HISTORY, dtype=torch.float32), + ) + + +def test_shared_selection_rejects_missing_row_and_width() -> None: + shared = _t(2, 8, 512, dtype=torch.int32) + with pytest.raises(AssertionError, match="no shared buffers"): + M.check_shared_selection(None, {3: 0}, 3, 512) + with pytest.raises(AssertionError, match="no row"): + M.check_shared_selection(shared, {4: 0}, 3, 512) + with pytest.raises(AssertionError, match="no row"): + M.check_shared_selection(shared, {3: 2}, 3, 512) + with pytest.raises(AssertionError, match="blocks per row"): + M.check_shared_selection(shared, {3: 1}, 3, 256) + + +def test_routing_admits_exactly_the_moe_backends_the_target_implements() -> None: + """`route` declines a backend request the constructor would assert on.""" + assert set(routing._MOE_BACKENDS) == set(M._TRUNK_MOE_BACKENDS) + + +def test_routing_admits_exactly_the_draft_lengths_the_target_implements() -> None: + """`route` declines a draft length the constructor or the first bind would assert on.""" + assert routing._MAX_DRAFT_LEN == M._MAX_VERIFY_WIDTH - 1 + assert routing._MAX_LEGACY_DRAFT_LEN == M._MAX_LEGACY_VERIFY_WIDTH - 1 + assert routing._MAX_CUTEDSL_DRAFT_LEN == M._CUTEDSL_MAX_DRAFT_LEN + + +def test_routing_admits_exactly_the_state_dtypes_the_target_serves() -> None: + assert routing._SSM_STATE_DTYPES == M._SSM_STATE_DTYPES diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_thop_attention_calls.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_thop_attention_calls.py new file mode 100644 index 000000000000..d9bbaa5c0ba0 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_thop_attention_calls.py @@ -0,0 +1,131 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Every `thop_attention` call of the qwen4_exp target names the KV storage it runs over. + +The attention op reads the pool through `quant_mode` and the two KV scales, and the three +have to agree with the pool's dtype. A call that leaves one at its default reads an fp8 pool +as bf16 (or the reverse) and still returns finite values, so no run notices. The check reads +`modeling.py` as text: it needs no device and no import of the target. + +The vision tower's call is the one exception: its attention is the encoder-only mode of the +op, which has no KV cache, so it passes `kv_scale_*` None and `quant_mode` 0 through +`_VISION_CALL_CONSTANTS`. +""" + +from __future__ import annotations + +import ast +from pathlib import Path + +import pytest + +import tensorrt_llm._torch._experimental.modeling_v2 as _modeling_v2 + +pytestmark = pytest.mark.cpu_only + +_MODELING = ( + Path(_modeling_v2.__file__).resolve().parent + / "models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py" +) +_KV_STORAGE_ARGS = ("quant_mode", "kv_scale_orig_quant", "kv_scale_quant_orig") +_TREE = ast.parse(_MODELING.read_text()) + + +def _thop_attention_calls() -> list[ast.Call]: + return [ + node + for node in ast.walk(_TREE) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "thop_attention" + ] + + +def _module_dict(name: str) -> ast.Call: + """The `dict(...)` a module-level `name = dict(...)` assigns.""" + for node in _TREE.body: + if isinstance(node, ast.Assign) and any( + isinstance(t, ast.Name) and t.id == name for t in node.targets + ): + assert isinstance(node.value, ast.Call), name + return node.value + raise AssertionError(f"modeling.py has no module-level {name}") + + +def _keywords(call: ast.Call) -> dict[str, ast.expr]: + return {kw.arg: kw.value for kw in call.keywords if kw.arg is not None} + + +def _unpacked(call: ast.Call) -> set[str]: + """Names of the `**name` arguments of `call`.""" + return { + kw.value.id for kw in call.keywords if kw.arg is None and isinstance(kw.value, ast.Name) + } + + +def _is_vision(call: ast.Call) -> bool: + return "_VISION_CALL_CONSTANTS" in _unpacked(call) + + +def test_the_target_calls_thop_attention_where_expected() -> None: + """Guards the checks below against going vacuous: the trunk and the draft layer each + call it from the prefill and the decode target, and the vision tower once.""" + calls = _thop_attention_calls() + assert len([c for c in calls if not _is_vision(c)]) >= 4 + assert len([c for c in calls if _is_vision(c)]) == 1 + + +_KV_CACHE_CALLS = [call for call in _thop_attention_calls() if not _is_vision(call)] + + +@pytest.mark.parametrize( + "call", _KV_CACHE_CALLS, ids=[f"call{i}" for i in range(len(_KV_CACHE_CALLS))] +) +def test_every_kv_cache_call_passes_the_pool_storage(call) -> None: + keywords = _keywords(call) + for name in _KV_STORAGE_ARGS: + assert name in keywords, f"thop_attention call at line {call.lineno} omits {name}" + value = keywords[name] + # Read off the core, whose `FP8` / `None` resolution all three follow. + assert isinstance(value, ast.Attribute) and value.attr == name, ( + f"line {call.lineno}: {name}={ast.unparse(value)}, not the core's `{name}`" + ) + + +def test_the_shared_constants_do_not_override_the_pool_storage() -> None: + """`**_CALL_CONSTANTS` is unpacked into every KV cache call; a storage argument in it + would either raise as a duplicate keyword or, if the call dropped its own, silently + pin the pool to one storage.""" + constants = {kw.arg for kw in _module_dict("_CALL_CONSTANTS").keywords} + assert constants + assert not constants & set(_KV_STORAGE_ARGS) + for call in _KV_CACHE_CALLS: + assert "_CALL_CONSTANTS" in _unpacked(call) + + +def test_the_step_state_carries_no_pool_storage() -> None: + """`**step` (each target's `step_args`) is unpacked into the same calls: a storage + argument there would collide with the call's own.""" + steps = [ + node + for node in ast.walk(_TREE) + if isinstance(node, ast.FunctionDef) and node.name == "step_args" + ] + assert len(steps) == 2, "one step_args per target (prefill, decode)" + for step in steps: + returned = [n for n in ast.walk(step) if isinstance(n, ast.Return)] + assert returned + for node in returned: + assert isinstance(node.value, ast.Call), "step_args returns dict(...)" + assert not {kw.arg for kw in node.value.keywords} & set(_KV_STORAGE_ARGS) + + +def test_the_vision_tower_runs_without_kv_storage() -> None: + """The encoder-only call states a bf16, unscaled, cache-less configuration.""" + constants = {kw.arg: kw.value for kw in _module_dict("_VISION_CALL_CONSTANTS").keywords} + assert ast.literal_eval(constants["quant_mode"]) == 0 + assert ast.literal_eval(constants["kv_scale_orig_quant"]) is None + assert ast.literal_eval(constants["kv_scale_quant_orig"]) is None + for call in _thop_attention_calls(): + if _is_vision(call): + assert not set(_keywords(call)) & set(_KV_STORAGE_ARGS) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_validate_request.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_validate_request.py new file mode 100644 index 000000000000..396a99268b1c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_validate_request.py @@ -0,0 +1,226 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""qwen4_exp's admission check fails a bad media request before the forward would assert. + +The check reads only the request's host-side payload, so these run on fake request +objects with meta pixel tensors: no GPU, vision ops or engine. Every refusal is a +`ValueError`, which the executor turns into that one request's error. +""" + +from __future__ import annotations + +import types +from types import SimpleNamespace + +import pytest +import torch + +__extra_import_path__ = [".."] # noqa: F841 -- repo's file-scoped import hook + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 + modeling as M, +) + +pytestmark = pytest.mark.cpu_only + +Shell = M.ModelingV2Qwen38FlashNextNvfp4Sm103Tp1 + +MERGE = 2 +PATCH = 16 +PIXEL_WIDTH = 4 +TEXT_TOKENS = 5 + + +def _shell(): + shell = SimpleNamespace( + model=SimpleNamespace(v_merge=MERGE, v_patch=PATCH, v_pixel_width=PIXEL_WIDTH) + ) + shell.validate_request = types.MethodType(Shell.validate_request, shell) + shell._media_grid = types.MethodType(Shell._media_grid, shell) + shell._check_item_order = Shell._check_item_order + return shell + + +def _media(modality, grids): + """A payload whose pixel tensor has the right shape and no storage.""" + patches = sum(t * h * w for t, h, w in grids) + key = "pixel_values" if modality == "image" else "pixel_values_videos" + grid_key = "image_grid_thw" if modality == "image" else "video_grid_thw" + return { + key: torch.empty(patches, PIXEL_WIDTH, device="meta"), + grid_key: torch.tensor(grids), + } + + +def _request(image=None, video=None, order=None, **overrides): + """A request whose prompt is `TEXT_TOKENS` text tokens around the placeholders the + media produce, with the position ids and placeholder cumsum the processor emits.""" + rows = sum(t * h * w for t, h, w in (image or []) + (video or [])) // (MERGE * MERGE) + prompt_len = rows + TEXT_TOKENS + data = { + "mrope_config": {"mrope_position_ids": torch.zeros(3, 1, prompt_len, dtype=torch.int32)}, + "multimodal_embed_mask_cumsum": torch.cat( + [torch.arange(1, rows + 1), torch.full((TEXT_TOKENS,), rows)] + ), + } + if image is not None: + data["image"] = _media("image", image) + if video is not None: + data["video"] = _media("video", video) + fields = {"py_multimodal_data": data, "py_mm_item_order": order, "prompt_len": prompt_len} + fields.update(overrides) + return SimpleNamespace(**fields) + + +def _admit(request): + _shell().validate_request(request) + + +def _item(modality, index): + return {"modality": modality, "index": index} + + +IMAGES = [(1, 4, 4), (1, 2, 2)] +VIDEOS = [(2, 2, 4), (2, 4, 4)] +MIXED_ORDER = [_item("video", 0), _item("image", 0), _item("video", 1), _item("image", 1)] + + +def test_clean_requests_are_admitted(): + _admit(_request(image=IMAGES)) + _admit(_request(video=VIDEOS)) + _admit(_request(image=IMAGES, video=VIDEOS, order=MIXED_ORDER)) + + +def test_a_request_without_media_costs_nothing(): + # Not even the other request attributes are read. + _admit(SimpleNamespace(py_multimodal_data=None)) + _admit(SimpleNamespace(py_multimodal_data={})) + # A warmup request carries position ids and no media. + _admit( + SimpleNamespace( + py_multimodal_data={"mrope_config": {"mrope_position_ids": torch.zeros(3, 1, 8)}} + ) + ) + # Precomputed embeddings arrive without pixels; the splice skips them. + _admit(SimpleNamespace(py_multimodal_data={"multimodal_embedding": torch.zeros(4, 3)})) + + +def test_a_video_past_the_per_call_row_count_is_admitted(): + # 33 temporal groups of 32x32 patches: 33792 rows in all, 1024 past one call. + _admit(_request(video=[(33, 32, 32)])) + + +def test_the_longest_admitted_frame_is_the_calls_ceiling(monkeypatch): + monkeypatch.setattr(M, "_VISION_CALL_MAX_PATCHES", 64) + _admit(_request(image=[(1, 8, 8)])) + _admit(_request(video=[(3, 4, 16)])) + with pytest.raises(ValueError, match="frames of 8x10 = 80 patches"): + _admit(_request(image=[(1, 8, 8), (1, 8, 10)])) + + +@pytest.mark.parametrize("modality", ["image", "video"]) +def test_a_frame_past_the_ceiling_fails_the_request(modality): + # 3072x3072 pixels: a 192x192 grid, 36864 patches in one frame. + request = _request(**{modality: [(1, 192, 192)]}) + with pytest.raises(ValueError, match=rf"{modality} 0 has frames of 192x192 = 36864 patches"): + _admit(request) + # The message names the pixels and the bound. + with pytest.raises(ValueError, match=r"\(3072x3072 pixels\).*at most 32768 patches"): + _admit(request) + + +def test_the_frame_bound_is_per_frame_not_per_item(): + # 4 frames of 128x128 patches (16384 each) are fine; 192x192 in the second item is not. + _admit(_request(video=[(4, 128, 128)])) + with pytest.raises(ValueError, match="video 1 has frames of 192x192"): + _admit(_request(video=[(4, 128, 128), (1, 192, 192)])) + + +@pytest.mark.parametrize( + "grid", [[(0, 4, 4)], [(1, 4, 5)], [(1, 3, 4)], [(1, 0, 4)]], ids=lambda g: str(g) +) +def test_a_malformed_grid_is_refused(grid): + request = _request(image=[(1, 4, 4)]) + request.py_multimodal_data["image"]["image_grid_thw"] = torch.tensor(grid) + with pytest.raises(ValueError, match="spatial_merge_size 2"): + _admit(request) + + +def test_a_grid_that_is_not_one_row_per_item_is_refused(): + request = _request(image=IMAGES) + request.py_multimodal_data["image"]["image_grid_thw"] = torch.tensor([1, 4, 4]) + with pytest.raises(ValueError, match="one \\(t, h, w\\) row per image"): + _admit(request) + + +def test_pixels_that_disagree_with_the_grid_are_refused(): + request = _request(video=VIDEOS) + video = request.py_multimodal_data["video"] + video["pixel_values_videos"] = torch.empty(5, PIXEL_WIDTH, device="meta") + with pytest.raises(ValueError, match="pixel_values_videos is .*describe 48 patch rows"): + _admit(request) + + +def test_a_payload_without_its_pixels_or_grid_is_refused(): + request = _request(video=VIDEOS) + del request.py_multimodal_data["video"]["video_grid_thw"] + with pytest.raises(ValueError, match="this target reads 'pixel_values_videos'"): + _admit(request) + + +def test_missing_position_ids_fail_the_request(): + request = _request(video=VIDEOS) + del request.py_multimodal_data["mrope_config"] + with pytest.raises(ValueError, match="mrope_position_ids.*got None"): + _admit(request) + + +@pytest.mark.parametrize("shape", [(3, 1, 4), (3, 1, 40), (2, 1, 12 + TEXT_TOKENS), (3, 12 + 5)]) +def test_mis_sized_position_ids_fail_the_request(shape): + request = _request(video=VIDEOS) + request.py_multimodal_data["mrope_config"]["mrope_position_ids"] = torch.zeros(*shape) + with pytest.raises(ValueError, match="one position per prompt token"): + _admit(request) + + +def test_a_mixed_request_needs_its_manifest(): + with pytest.raises(ValueError, match="needs mm_item_order"): + _admit(_request(image=IMAGES, video=VIDEOS)) + with pytest.raises(ValueError, match="needs mm_item_order"): + _admit(_request(image=IMAGES, video=VIDEOS, order=[])) + + +def test_a_manifest_that_disagrees_with_the_payload_is_refused(): + with pytest.raises(ValueError, match="lists image items \\[0\\]; the payload carries 2"): + _admit(_request(image=IMAGES, video=VIDEOS, order=MIXED_ORDER[:3])) + with pytest.raises(ValueError, match="lists video items \\[0, 1, 0\\]"): + _admit(_request(image=IMAGES, video=VIDEOS, order=MIXED_ORDER + [_item("video", 0)])) + with pytest.raises(ValueError, match="is not an item of this request"): + _admit(_request(image=IMAGES, video=VIDEOS, order=MIXED_ORDER + [_item("audio", 0)])) + with pytest.raises(ValueError, match="is not an item of this request"): + _admit(_request(image=IMAGES, video=VIDEOS, order=[{"modality": "image"}])) + + +def test_a_single_modality_request_needs_no_manifest(): + _admit(_request(image=IMAGES, order=None)) + _admit(_request(video=VIDEOS, order=None)) + + +def test_a_prompt_whose_placeholders_disagree_with_the_media_is_refused(): + request = _request(video=VIDEOS) + request.py_multimodal_data["multimodal_embed_mask_cumsum"] = torch.arange(1, 9) + with pytest.raises(ValueError, match="holds 8 image/video placeholder tokens.*12 embedding"): + _admit(request) + + +def test_the_placeholder_count_is_the_sum_over_both_modalities(): + request = _request(image=IMAGES, video=VIDEOS, order=MIXED_ORDER) + request.py_multimodal_data["multimodal_embed_mask_cumsum"] = torch.arange(1, 14) + with pytest.raises(ValueError, match="holds 13 image/video placeholder tokens.*17"): + _admit(request) + + +def test_a_request_without_a_cumsum_skips_the_count(): + request = _request(video=VIDEOS) + del request.py_multimodal_data["multimodal_embed_mask_cumsum"] + _admit(request) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_weight_table.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_weight_table.py index bdb99113b301..4145ef96632c 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_weight_table.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_weight_table.py @@ -97,9 +97,16 @@ def test_mtp_rows_appear_only_when_mtp_is_live(): on = W.MODEL_WEIGHTS.manifest(_core(True)) added = set(on) - set(off) assert set(off) <= set(on) - assert added and all(k.startswith(f"l{_TEXT['num_hidden_layers']}_mtp_") for k in added) - # The draft layer's own QSA indexer is never loaded; see weights.py. - assert not any("indexer" in str(src) for k in added for src, _, _ in on[k]) + m = f"l{_TEXT['num_hidden_layers']}_mtp_" + assert added and all(k.startswith(m) for k in added) + # The draft layer runs QSA over its own indexer, read from the module's keys. + indexer = {f"{m}idx_w", f"{m}idx_qn", f"{m}idx_kn"} + assert indexer <= added + assert all( + str(src).startswith("mtp.layers.0.self_attn.indexer.") + for k in indexer + for src, _, _ in on[k] + ) def _reference_declaration(mtp: bool) -> dict: @@ -223,6 +230,9 @@ def shared_expert(prefix): hyper_connection(f"{m}ahc") hyper_connection(f"{m}mhc") gated_attention(m) + P(f"{m}idx_w", index_width, hidden) + P(f"{m}idx_qn", t.indexer_head_dim) + P(f"{m}idx_kn", t.indexer_head_dim) P(f"{m}router", experts, hidden) P(f"{m}fc1_w", experts, 2 * inter, hidden, dtype=f8) P(f"{m}fc1_s", experts, (2 * inter) // 128, hidden // 128, dtype=f32) @@ -248,3 +258,30 @@ def test_the_table_declares_what_the_init_block_did(mtp): ) wrong = {k: (got[k], want[k]) for k in got if got[k] != want[k]} assert not wrong, list(wrong.items())[:8] + + +def test_cutedsl_fc1_pairs_up_and_gate_rows_in_64_row_groups(): + """Destination rows `[128k, 128k + 64)` hold up rows `[64k, 64k + 64)` and the next 64 + hold the matching gate rows, which sit `I` rows later in the `[up ; gate]` stack.""" + inter = _TEXT["moe_intermediate_size"] + core = SimpleNamespace(trunk_moe_backend="CUTEDSL") + perm = W._fc1_perm(2 * inter, "cpu", core) + assert sorted(perm.tolist()) == list(range(2 * inter)) + for k in range(inter // 64): + assert perm[128 * k : 128 * k + 64].tolist() == list(range(64 * k, 64 * k + 64)) + assert perm[128 * k + 64 : 128 * k + 128].tolist() == list( + range(inter + 64 * k, inter + 64 * k + 64) + ) + + +def test_cutedsl_fc2_keeps_the_checkpoint_row_order_and_trtllm_gen_shuffles_it(): + """trtllm-gen's 32-row block shuffle moves source row `4u + v` to destination row + `8v + u`; CuteDSL reads `down_proj` as the checkpoint stores it.""" + rows = torch.arange(64).unsqueeze(1).expand(64, 8).contiguous() + kept = W._expert_fc2_w(rows, SimpleNamespace(trunk_moe_backend="CUTEDSL")) + assert torch.equal(kept, rows) + shuffled = W._expert_fc2_w(rows, SimpleNamespace(trunk_moe_backend="TRTLLM")) + for u in range(8): + for v in range(4): + assert int(shuffled[8 * v + u, 0]) == 4 * u + v + assert int(shuffled[32 + 8 * v + u, 0]) == 32 + 4 * u + v From c8ed8cb05d473b9ab13c3df7775d519aa25bc9ea Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Sun, 18 Oct 2026 22:59:57 -0700 Subject: [PATCH 03/28] [TRTLLM-17033][test] Accuracy gates and test lists for the Qwen3.8-Flash-Next modeling_v2 features Integration gates in `tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py` (MTP 3/5/7, CuteDSL + MTP3, long context, the legacy GDN verify path, video, bf16 GDN state, FP8 and NVFP4 KV, the breakable prefill CUDA graph), their gsm8k and acceptance-length references, and the l0_b300 / l0_cpu test-list entries. Folds these commits of the review branch (their changes to the files of this commit): - Run the Qwen3.8-Flash-Next trunk routed experts on CUTEDSL (5405a3d6ee) - Complete MTP for the Qwen3.8-Flash-Next modeling_v2 target (0a1b9c48da) - Store the Qwen3.8-Flash-Next KV cache as FP8 when kv_cache_quant_algo is FP8 (744607bf5c) - Cover the qwen4_exp KV storage derivation and attention-call arguments on CPU (d64f96c637) - Check the CuteDSL expert weight path from checkpoint tensors to output (cd749d61d1) - Add MTP3, CuteDSL+MTP3, long-context, GDN-legacy and video gates for qwen4_exp (3a4161b655) - State the final measured run time of the modeling_v2 unit directory in l0_b300.yml (4e34d8bb52) - Gate the qwen4_exp target at five and seven MTP drafts (ffe8aafe54) - Gate the qwen4_exp target with the Gated-DeltaNet state in bf16 (97949eb931) - Reword the engine slot-count comments and restate the unit directory's run time (025a17c0ab) - Gate the qwen4_exp target under the breakable prefill CUDA graph (02df09606a) - Gate the qwen4_exp target over an NVFP4 KV cache (ef6918698a) Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../accuracy/modeling_v2/test_qwen4_exp.py | 992 +++++++++++++++++- .../references/acceptance_length.yaml | 147 ++- .../defs/accuracy/references/gsm8k.yaml | 12 + .../test_lists/test-db/l0_b300.yml | 43 +- .../integration/test_lists/test-db/l0_cpu.yml | 12 +- 5 files changed, 1156 insertions(+), 50 deletions(-) diff --git a/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py b/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py index f7764bd3fa60..f24a4574a41f 100644 --- a/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py +++ b/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py @@ -19,16 +19,35 @@ and report the built-in's numbers as the target's. """ +import random +import re +import time + import pytest +from PIL import Image, ImageDraw +from transformers import AutoProcessor, AutoTokenizer from tensorrt_llm import LLM from tensorrt_llm._torch._experimental.modeling_v2 import MODELING_V2_VALIDATE_ENV +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 + modeling as target_modeling, +) +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph_runner import BreakableCUDAGraphRunner from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.executor.utils import RequestError +from tensorrt_llm.inputs import VideoData +from tensorrt_llm.inputs.utils import ( + ConversationMessage, + MultimodalDataTracker, + apply_chat_template, + apply_mm_placeholders, +) from tensorrt_llm.llmapi import ( CudaGraphConfig, KvCacheConfig, MoeConfig, MTPDecodingConfig, + QSASparseAttentionConfig, SamplingParams, ) @@ -46,11 +65,6 @@ get_sm_version() != 103, reason="modeling_v2 targets in this batch are certified on sm_103 only" ) -# Speculative cases need the legacy Gated-DeltaNet verify path: the engine's -# default (replay) commits through a callable the target's closed op vocabulary -# has no entry for, and the target asserts the legacy path is selected. -_GDN_LEGACY = {"TRTLLM_USE_GDN_REPLAY": "0"} - # Check the target's step contract on the first forward, against metadata the # engine prepared; the flag has to reach the worker rank. _VALIDATE = {MODELING_V2_VALIDATE_ENV: "1"} @@ -67,6 +81,268 @@ def _require_target(monkeypatch: pytest.MonkeyPatch, **env: str) -> dict: return dict(modeling_v2="require", env_overrides=dict(env)) +def _target_core(llm: LLM): + """The target's core module. + + Reachable only with the executor in this process (``TLLM_WORKER_USE_SINGLE_PROCESS=1``), + where the iteration stats that ``compute_acceptance_length`` reads are not collected. + """ + return llm._executor.engine.model_engine.model.model + + +def _count_prefill_graph_replays(monkeypatch: pytest.MonkeyPatch) -> dict[int, int]: + """Replays of the breakable prefill graph by bucket, filled as the engine runs. + + Needs the executor in this process. Without it a case that sets + ``prefill_cuda_graph_backend="breakable"`` could pass on the eager path. + """ + replays: dict[int, int] = {} + replay = BreakableCUDAGraphRunner.replay + + def counting_replay(self, num_tokens): + replays[num_tokens] = replays.get(num_tokens, 0) + 1 + return replay(self, num_tokens) + + monkeypatch.setattr(BreakableCUDAGraphRunner, "replay", counting_replay) + return replays + + +def _acceptance_length_from_requests(outputs, draft_len: int) -> float: + """Mean acceptance length from the requests' own speculative-decoding metrics. + + Needs ``SamplingParams(return_perf_metrics=True)``. A draft loop yields one target token + plus the drafts it accepted. + """ + accepted = drafted = 0 + for output in outputs: + spec = output.outputs[0].request_perf_metrics.speculative_decoding + accepted += spec.total_accepted_draft_tokens + drafted += spec.total_draft_tokens + assert drafted > 0, "no request drafted a token" + return 1.0 + accepted * draft_len / drafted + + +# ---------------------------------------------------------------- long prompts --- +# Prompts past the QSA indexer's 2048-token budget, so every decode step of the +# target and of its MTP draft layer takes the sparse attention path. + +_PASSKEY_LENGTHS = (3000, 4000, 5000, 6000) +_PASSKEY_LENGTHS_LONG = (3000, 6000, 9000, 12000) +# 33 prompts, each past the indexer's budget, 32 of which one batch carries. +_FORCED_LENGTHS = tuple(range(2200, 6400, 400)) +_FORCED_TOKENS = 64 +_PASSKEY_DEPTHS = (0.15, 0.5, 0.85) +_PASSKEY_MAX_TOKENS = 160 + +_COMPONENTS = ( + "scheduler", + "allocator", + "tokenizer", + "router", + "profiler", + "exporter", + "planner", + "balancer", + "resolver", + "dispatcher", +) +_ACTIONS = ( + "reads", + "rewrites", + "caches", + "validates", + "forwards", + "compacts", + "audits", + "replays", + "indexes", + "batches", +) +_OBJECTS = ( + "pending requests", + "block tables", + "shard manifests", + "routing hints", + "retry budgets", + "lease records", + "checkpoint deltas", + "quota counters", +) +_CLAUSES = ( + "before the next flush", + "whenever the queue drains", + "after every shard reports", + "once the lease expires", + "while the window is open", + "until the barrier clears", +) +_STATES = ("consistent", "bounded", "ordered", "reproducible", "idle") +_TOPICS = ("rules", "limits", "guarantees", "hooks", "settings") + + +def _section(rng: random.Random, number: int) -> str: + heading = f"## Section {number}: {rng.choice(_COMPONENTS).capitalize()} {rng.choice(_TOPICS)}" + sentences = [ + f"The {rng.choice(_COMPONENTS)} {rng.choice(_ACTIONS)} the {rng.choice(_OBJECTS)} " + f"{rng.choice(_CLAUSES)}, so the {rng.choice(_OBJECTS)} stay {rng.choice(_STATES)}." + for _ in range(6) + ] + return heading + "\n\n" + " ".join(sentences) + + +def _passkey_prompts(tokenizer, lengths=_PASSKEY_LENGTHS) -> list[tuple[str, str]]: + """``(chat prompt, passkey)`` pairs over ``lengths`` x ``_PASSKEY_DEPTHS``. + + Each prompt is a generated document of about that many tokens with one pass key sentence + at that relative depth, the same bytes in every run. + """ + rng = random.Random(20261014) + question = ( + "Based on the document above, first state the secret pass key on a line of its own " + "in the form 'Pass key: '. Then list, one per line, the section headings " + "that appear in the document, in order." + ) + prompts = [] + for length in lengths: + for depth in _PASSKEY_DEPTHS: + key = str(rng.randint(100000, 999999)) + sections, tokens = [], 0 + while tokens < length - 150: + sections.append(_section(rng, len(sections) + 1)) + tokens += len(tokenizer.encode(sections[-1])) + needle = ( + f"IMPORTANT NOTE FOR THE READER: the secret pass key of this document is {key}. " + f"Remember the pass key {key} exactly." + ) + position = max(1, min(len(sections) - 1, round(depth * len(sections)))) + body = sections[:position] + [needle] + sections[position:] + text = "Document:\n\n" + "\n\n".join(body) + "\n\n" + question + prompt = tokenizer.apply_chat_template( + [{"role": "user", "content": text}], + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + prompts.append((prompt, key)) + return prompts + + +# ------------------------------------------------------------------ word problems --- + +_WORD_PROBLEM_SUFFIX = "Solve it step by step and finish with 'The answer is .'" +_WORD_PROBLEMS = ( + ( + "A bakery bakes 7 trays of 12 muffins and sells all but 9. How many muffins did it sell?", + "75", + ), + ( + "Tom has 3 boxes of 24 pencils and gives away 19 pencils. How many pencils does he " + "have left?", + "53", + ), + ( + "A train travels 60 miles per hour for 2.5 hours, then 40 miles per hour for 3 hours. " + "How many miles does it travel in total?", + "270", + ), + ( + "Sara buys 5 notebooks at $3 each and 4 pens at $2 each and pays with a $50 bill. " + "How many dollars of change does she get?", + "27", + ), + ("A rectangle is 14 cm long and 9 cm wide. What is its perimeter in cm?", "46"), + ( + "125 students attend; 40% play soccer and one fifth of the others play tennis. " + "How many students play tennis?", + "15", + ), +) + + +# ---------------------------------------------------------------- media inputs --- + +_MODEL_TYPE = "qwen4_exp" +_RED, _BLUE, _GREEN = (220, 30, 30), (30, 60, 220), (30, 170, 60) +# name -> (square colour, direction, colour word, direction phrases an answer may use) +_VIDEOS = { + "red_rightward": (_RED, "right", "red", ("left to right", "to the right", "rightward")), + "blue_leftward": (_BLUE, "left", "blue", ("right to left", "to the left", "leftward")), + "green_downward": (_GREEN, "down", "green", ("top to bottom", "downward", "down")), +} +_VIDEO_QUESTION = ( + "What colour is the moving square, and in which direction does it move across the frame? " + "Answer in one short sentence." +) +_MIXED_QUESTION = ( + "The image shows a square and the video shows a moving square. Answer in exactly this " + "format: 'Image: . Video: .'" +) + + +def _synthetic_video(name: str, num_frames: int = 16, size: int = 256) -> VideoData: + """A square of the named colour crossing a plain background in the named direction. + + The frames are built in memory, so no video decoder is involved. + """ + color, direction, _, _ = _VIDEOS[name] + side, margin = 48, 8 + frames = [] + for i in range(num_frames): + t = i / (num_frames - 1) + travel = margin + t * (size - side - 2 * margin) + centre = (size - side) // 2 + x, y = { + "right": (travel, centre), + "left": (size - side - travel, centre), + "down": (centre, travel), + }[direction] + frame = Image.new("RGB", (size, size), (235, 235, 235)) + ImageDraw.Draw(frame).rectangle((x, y, x + side, y + side), fill=color) + frames.append(frame) + return VideoData( + frames=frames, + metadata=dict( + total_num_frames=num_frames, + fps=2.0, + duration=num_frames / 2.0, + frames_indices=list(range(num_frames)), + ), + ) + + +def _synthetic_image(color: tuple, size: int = 256) -> Image.Image: + """A square of ``color`` on a plain background.""" + image = Image.new("RGB", (size, size), (235, 235, 235)) + quarter = size // 4 + ImageDraw.Draw(image).rectangle((quarter, quarter, size - quarter, size - quarter), fill=color) + return image + + +def _media_request(tokenizer, processor, text: str, items: list[tuple[str, object]]) -> dict: + """One request carrying ``items``, ``(modality, data)`` pairs in prompt order. + + It holds the chat prompt with its placeholders, the media and the ``mm_item_order`` + manifest that splices several modalities back in prompt order. + """ + tracker = MultimodalDataTracker(_MODEL_TYPE) + for modality, data in items: + tracker.add_data(modality, data) + counts = tracker.placeholder_counts() + message = ConversationMessage(role="user", content=text, media=[]) + apply_mm_placeholders(_MODEL_TYPE, message, counts, tracker) + prompt = apply_chat_template( + model_type=_MODEL_TYPE, + tokenizer=tokenizer, + processor=processor, + conversation=[message], + add_generation_prompt=True, + mm_placeholder_counts=[counts], + chat_template_kwargs=dict(enable_thinking=False), + ) + media, _ = tracker.retrieve_all_sync() + return dict(prompt=prompt, multi_modal_data=media, mm_item_order=tracker.item_order()) + + class TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1(LlmapiAccuracyTestHarness): """qwen3.8-flash-next-nvfp4 / sm_103 / tp1: text, MTP and multimodal.""" @@ -88,13 +364,39 @@ class TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1(LlmapiAccuracyTestHarness): GSM8K_MAX_OUTPUT_LEN = 512 # `mamba_ssm_cache_dtype` is float32 because the checkpoint declares - # `mamba_ssm_dtype: float32` (the engine default is bfloat16). `dtype` is - # left at the checkpoint's bf16 KV, which is what the anchors were measured at. + # `mamba_ssm_dtype: float32` (the engine default is bfloat16); the bf16-state cases + # set it to bfloat16. `dtype` is left at the checkpoint's bf16 KV, which is what the + # anchors were measured at; only the fp8 and NVFP4 KV cases set it. KV_CACHE = KvCacheConfig( free_gpu_memory_fraction=0.6, enable_block_reuse=True, mamba_ssm_cache_dtype="float32", ) + KV_CACHE_BF16_STATE = KvCacheConfig( + free_gpu_memory_fraction=0.6, + enable_block_reuse=True, + mamba_ssm_cache_dtype="bfloat16", + ) + KV_CACHE_FP8 = KvCacheConfig( + dtype="fp8", + free_gpu_memory_fraction=0.6, + enable_block_reuse=True, + mamba_ssm_cache_dtype="float32", + ) + # Block reuse is on, which (with chunked prefill and the MTP drafts) is what lets the + # attention op read an NVFP4 pool: its kernels exist for the paged-context path only. + KV_CACHE_NVFP4 = KvCacheConfig( + dtype="nvfp4", + free_gpu_memory_fraction=0.6, + enable_block_reuse=True, + mamba_ssm_cache_dtype="float32", + ) + # Block reuse is off for the vision path, which requires it. + KV_CACHE_MULTIMODAL = KvCacheConfig( + free_gpu_memory_fraction=0.6, + enable_block_reuse=False, + mamba_ssm_cache_dtype="float32", + ) # Required, not a tuning choice: the MTP draft layer's FP8 block-scale # experts are not supported by the default CUTLASS MoE on sm_103. Set for @@ -110,11 +412,64 @@ class TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1(LlmapiAccuracyTestHarness): moe_config=MOE, ) - # One-model MTP-Eagle with the checkpoint's single MTP module. + # The breakable prefill CUDA graph is opt-in. The buckets cover the gsm8k prompts; with the + # 2048-token budget every chunk of a longer prompt lands on one as well. + BREAKABLE_PREFILL = dict( + prefill_cuda_graph_backend="breakable", + prefill_capture_num_tokens=[128, 256, 512, 1024, 2048, 3072], + ) + ENGINE_BREAKABLE = dict(ENGINE, **BREAKABLE_PREFILL) + ENGINE_BREAKABLE_CHUNKED = dict(ENGINE_BREAKABLE, max_num_tokens=2048) + + # One-model MTP-Eagle with the checkpoint's single MTP module replayed + # `max_draft_len` times. MTP1 = MTPDecodingConfig(max_draft_len=1) + MTP3 = MTPDecodingConfig(max_draft_len=3) + MTP5 = MTPDecodingConfig(max_draft_len=5) + MTP7 = MTPDecodingConfig(max_draft_len=7) - # Acceptance-length anchor, keyed by target and variant. + # Acceptance-length anchors in references/acceptance_length.yaml, keyed by + # target and variant. ACCEPTANCE_KEY = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp1" + ACCEPTANCE_KEY_FP8_KV = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp1_fp8_kv" + ACCEPTANCE_KEY_NVFP4_KV = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp1_nvfp4_kv" + ACCEPTANCE_KEY_MTP3 = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3" + ACCEPTANCE_KEY_MTP3_CUTEDSL = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_cutedsl" + ACCEPTANCE_KEY_MTP5 = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp5" + ACCEPTANCE_KEY_MTP7 = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp7" + ACCEPTANCE_KEY_MTP7_LONG_CONTEXT = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp7_long_context" + ACCEPTANCE_KEY_MTP3_BF16_STATE = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_bf16_state" + ACCEPTANCE_KEY_MTP7_BF16_STATE_LONG_CONTEXT = ( + "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp7_bf16_state_long_context" + ) + ACCEPTANCE_KEY_LONG_CONTEXT = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_long_context" + ACCEPTANCE_KEY_FP8_KV_LONG_CONTEXT = ( + "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_fp8_kv_long_context" + ) + ACCEPTANCE_KEY_NVFP4_KV_LONG_CONTEXT = ( + "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_nvfp4_kv_long_context" + ) + ACCEPTANCE_KEY_GDN_LEGACY = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp1_gdn_legacy" + + def _gsm8k_with_acceptance(self, mocker, monkeypatch, engine, speculative_config, key, label): + """gsm8k against the reference row the engine resolves, then the acceptance length. + + The acceptance length of the same run is held to the floor under ``key``. + """ + mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", self.GSM8K_MAX_OUTPUT_LEN) + + with LLM( + self.MODEL_PATH, + speculative_config=speculative_config, + enable_iter_perf_stats=True, + **engine, + **_require_target(monkeypatch, **_VALIDATE), + ) as llm: + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) + acceptance_length = compute_acceptance_length(llm) + print(f"[AL] {label} acceptance_length = {acceptance_length:.3f}") + assert_acceptance_length(key, acceptance_length) @skip_not_sm103 def test_gsm8k(self, mocker, monkeypatch): @@ -129,6 +484,62 @@ def test_gsm8k(self, mocker, monkeypatch): task = GSM8K(self.MODEL_NAME) task.evaluate(llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) + @skip_not_sm103 + def test_gsm8k_breakable_prefill_cuda_graph(self, mocker, monkeypatch): + """The text model with ``prefill_cuda_graph_backend="breakable"``, against the same gsm8k anchor. + + Every prefill step of up to 3072 tokens is padded to a captured bucket and replayed, so + the padding rows, the four eager breaks and the per-bucket graphs run under real + prompts. That replays happen is shown by ``test_breakable_prefill_matches_eager``. + """ + mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", self.GSM8K_MAX_OUTPUT_LEN) + + with LLM( + self.MODEL_PATH, + **self.ENGINE_BREAKABLE, + **_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_gsm8k_cutedsl_moe(self, mocker, monkeypatch): + """The text model with the trunk's routed experts on CuteDSL, against the same gsm8k anchor. + + Same as ``test_gsm8k`` with ``moe_config.backend="CUTEDSL"``. + """ + mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", self.GSM8K_MAX_OUTPUT_LEN) + engine = dict(self.ENGINE, moe_config=MoeConfig(backend="CUTEDSL")) + + with LLM( + self.MODEL_PATH, + **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_gsm8k_bf16_ssm_state(self, mocker, monkeypatch): + """The text model with the recurrent state stored in bf16, against the fp32 gsm8k anchor. + + ``mamba_ssm_cache_dtype="bfloat16"`` (what ``auto`` resolves to for this checkpoint) rounds + the state to bf16 on every decode token and every prefill call, a lossy substitution + for the checkpoint's fp32 state. It reads the same reference row as ``test_gsm8k``, so + the rounding has to stay inside that row's tolerance. The decode step runs its + recurrence in the batch-invariant form that a bf16 pool needs. + """ + mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", self.GSM8K_MAX_OUTPUT_LEN) + engine = dict(self.ENGINE, kv_cache_config=self.KV_CACHE_BF16_STATE) + + with LLM( + self.MODEL_PATH, + **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. @@ -136,37 +547,511 @@ def test_mtp1_gsm8k_and_acceptance(self, mocker, monkeypatch): 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. + tripwire. """ - mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", self.GSM8K_MAX_OUTPUT_LEN) + self._gsm8k_with_acceptance( + mocker, monkeypatch, self.ENGINE, self.MTP1, self.ACCEPTANCE_KEY, "mtp1" + ) + + @skip_not_sm103 + def test_mtp1_fp8_kv_gsm8k_and_acceptance(self, mocker, monkeypatch): + """MTP with one draft over an fp8 KV cache: the FP8 KV row of references/gsm8k.yaml. + + ``kv_cache_config.dtype="fp8"`` makes the engine resolve ``kv_cache_quant_algo`` FP8, + which the reference lookup keys on: this case reads the FP8 row where + ``test_mtp1_gsm8k_and_acceptance`` (no KV quantization) reads the bf16 MTP row. The + draft layer runs over the same fp8 pool; its acceptance length has a floor of its own, + since fp8 storage can cost drafts the bf16 cache would have kept. + """ + engine = dict(self.ENGINE, kv_cache_config=self.KV_CACHE_FP8) + self._gsm8k_with_acceptance( + mocker, monkeypatch, engine, self.MTP1, self.ACCEPTANCE_KEY_FP8_KV, "mtp1 fp8 kv" + ) + + @skip_not_sm103 + def test_mtp1_nvfp4_kv_gsm8k_and_acceptance(self, mocker, monkeypatch): + """MTP with one draft over an NVFP4 KV cache: the NVFP4 KV row of references/gsm8k.yaml. + + ``kv_cache_config.dtype="nvfp4"`` makes the engine resolve ``kv_cache_quant_algo`` + NVFP4, which the reference lookup keys on. A lossy substitution of the checkpoint's bf16 + K/V (e2m1 values under e4m3 block scales), so the accuracy and the acceptance length + each have a row of their own. The draft layer runs over the same NVFP4 pool. + """ + engine = dict(self.ENGINE, kv_cache_config=self.KV_CACHE_NVFP4) + self._gsm8k_with_acceptance( + mocker, monkeypatch, engine, self.MTP1, self.ACCEPTANCE_KEY_NVFP4_KV, "mtp1 nvfp4 kv" + ) + + @skip_not_sm103 + def test_mtp3_gsm8k_and_acceptance(self, mocker, monkeypatch): + """MTP with three drafts: the MTP row of references/gsm8k.yaml, and a length floor. + + Three drafts replay the one MTP module three times, with the verify step running four + rows per request, a different shape from ``test_mtp1_gsm8k_and_acceptance`` on every + attention and recurrent-state call of the verify step. Drafts after the first can be + rejected while the first stays accepted, which holds the length near the one-draft + value, so the floor has its own entry. + """ + self._gsm8k_with_acceptance( + mocker, monkeypatch, self.ENGINE, self.MTP3, self.ACCEPTANCE_KEY_MTP3, "mtp3" + ) + + @skip_not_sm103 + def test_mtp3_breakable_prefill_gsm8k_and_acceptance(self, mocker, monkeypatch): + """MTP with three drafts over replayed prefill steps: the MTP row, and the length floor. + + The padded rows must never reach the draft loop: the worker takes the batch's own rows + and the draft layer's first step runs over them eagerly. A draft layer fed a wrong + hidden state loses acceptance length without changing accuracy. + """ + self._gsm8k_with_acceptance( + mocker, + monkeypatch, + self.ENGINE_BREAKABLE, + self.MTP3, + self.ACCEPTANCE_KEY_MTP3, + "mtp3 breakable prefill", + ) + + @skip_not_sm103 + def test_mtp3_cutedsl_gsm8k_and_acceptance(self, mocker, monkeypatch): + """MTP with three drafts over a trunk whose routed experts run on CuteDSL. + + The trunk's CuteDSL experts and the draft layer's FP8 block-scale experts run in one + forward only here: ``test_gsm8k_cutedsl_moe`` has no draft layer and the other MTP + cases run the trunk on trtllm-gen. Acceptance length detects a trunk hidden state the + draft layer cannot use, which accuracy does not. + """ + engine = dict(self.ENGINE, moe_config=MoeConfig(backend="CUTEDSL")) + self._gsm8k_with_acceptance( + mocker, + monkeypatch, + engine, + self.MTP3, + self.ACCEPTANCE_KEY_MTP3_CUTEDSL, + "mtp3 cutedsl", + ) + + @skip_not_sm103 + def test_mtp5_gsm8k_and_acceptance(self, mocker, monkeypatch): + """MTP with five drafts: the MTP row of references/gsm8k.yaml, and a length floor. + + Six verify rows per request, a width that no other case runs: the cached replay at + ``T = 6`` and ``thop_attention`` at ``predicted_tokens_per_seq`` 6, with the all-layer + commit from 16 requests on. + """ + self._gsm8k_with_acceptance( + mocker, monkeypatch, self.ENGINE, self.MTP5, self.ACCEPTANCE_KEY_MTP5, "mtp5" + ) + + @skip_not_sm103 + def test_mtp7_gsm8k_and_acceptance(self, mocker, monkeypatch): + """MTP with seven drafts, the longest the target runs: the MTP row and a length floor. + + Eight verify rows per request is the cached replay kernel's ceiling (``T = 8``, history + 16): a slot that accepts every draft fills its history every other step, so the + checkpoint commit, in-kernel below 16 requests and all-layer from there on, runs twice + as often as at three drafts. The engine captures a CUDA graph for every batch size up + to 32, so every captured step carries up to 256 verify rows. + """ + self._gsm8k_with_acceptance( + mocker, monkeypatch, self.ENGINE, self.MTP7, self.ACCEPTANCE_KEY_MTP7, "mtp7" + ) + + @skip_not_sm103 + def test_mtp3_bf16_ssm_state_gsm8k_and_acceptance(self, mocker, monkeypatch): + """MTP with three drafts over a bf16 recurrent state: the MTP row, and a length floor. + + The cached replay rebuilds the state from a checkpoint that is now bf16 and rounds it + once at every commit; drafts after the first are the ones a degraded state would lose, + so the length floor has its own entry. + """ + engine = dict(self.ENGINE, kv_cache_config=self.KV_CACHE_BF16_STATE) + self._gsm8k_with_acceptance( + mocker, + monkeypatch, + engine, + self.MTP3, + self.ACCEPTANCE_KEY_MTP3_BF16_STATE, + "mtp3 bf16 state", + ) + + def _long_context_passkey( + self, monkeypatch, speculative_config, engine, index_share, lengths, key_name, label + ): + """Pass key retrieval over prompts past the indexer's budget, at ``speculative_config``. + + Runs with the executor in this process to read the draft layer's op calls; the + acceptance length comes from the requests' own metrics. Two identical requests in one + batch must produce the same tokens, which fails if a request's shared selection leaks + into another slot. + """ + monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") + reuse_calls = [] + reuse = target_modeling.qsa_mtp_reuse_topk + + def counting_reuse(*args, **kwargs): + reuse_calls.append(1) + return reuse(*args, **kwargs) + + monkeypatch.setattr(target_modeling, "qsa_mtp_reuse_topk", counting_reuse) + engine = dict( + engine, + sparse_attention_config=QSASparseAttentionConfig( + index_share_for_mtp_iteration=index_share + ), + ) + tokenizer = AutoTokenizer.from_pretrained(self.MODEL_PATH) + prompts = _passkey_prompts(tokenizer, lengths) + replays = _count_prefill_graph_replays(monkeypatch) + # One prompt twice, far apart in the batch, so the two copies hold different slots. + texts = [prompt for prompt, _ in prompts] + [prompts[len(prompts) // 2][0]] + + with LLM( + self.MODEL_PATH, + speculative_config=speculative_config, + **engine, + **_require_target(monkeypatch, **_VALIDATE), + ) as llm: + assert _target_core(llm).mtp_index_share == index_share + outputs = llm.generate( + texts, + SamplingParams( + max_tokens=_PASSKEY_MAX_TOKENS, temperature=0, return_perf_metrics=True + ), + use_tqdm=False, + ) + + missed = [ + key + for (_, key), output in zip(prompts, outputs) + if key not in output.outputs[0].text[:_PASSKEY_MAX_TOKENS] + ] + assert not missed, f"pass keys not stated back: {missed} of {[key for _, key in prompts]}" + twins = outputs[len(prompts) // 2], outputs[-1] + assert twins[0].outputs[0].token_ids == twins[1].outputs[0].token_ids + if engine.get("prefill_cuda_graph_backend") == "breakable": + # Every prompt is longer than the chunk budget, so its chunks fill the top bucket. + assert replays.get(2048, 0) >= len(texts), replays + else: + assert not replays, replays + if index_share: + assert reuse_calls, "IndexShare is on and the reuse op never ran" + else: + assert not reuse_calls, "IndexShare is off and the reuse op ran" + acceptance_length = _acceptance_length_from_requests( + outputs, speculative_config.max_draft_len + ) + print(f"[AL] {label} acceptance_length = {acceptance_length:.3f}") + assert_acceptance_length(getattr(self, key_name), acceptance_length) + + @pytest.mark.parametrize( + "kv_cache, index_share, breakable, key_name", + [ + pytest.param("KV_CACHE", True, False, "ACCEPTANCE_KEY_LONG_CONTEXT", id="index_share"), + pytest.param( + "KV_CACHE_FP8", False, False, "ACCEPTANCE_KEY_FP8_KV_LONG_CONTEXT", id="fp8_kv" + ), + pytest.param( + "KV_CACHE_NVFP4", + False, + False, + "ACCEPTANCE_KEY_NVFP4_KV_LONG_CONTEXT", + id="nvfp4_kv", + ), + pytest.param( + "KV_CACHE", True, True, "ACCEPTANCE_KEY_LONG_CONTEXT", id="breakable_prefill" + ), + ], + ) + @skip_not_sm103 + def test_mtp3_long_context_passkey( + self, monkeypatch, kv_cache, index_share, breakable, key_name + ): + """Pass key retrieval over 3k-6k-token prompts, where the sparse attention path runs. + + The gsm8k cases never leave the dense path (their prompts are under the indexer's + 2048-token budget), so this is the only case that runs the target's QSA selection, the + draft layer's selection and, over an fp8 or NVFP4 pool, the sparse store and GQA kernels + of that format (the draft layer's included) end to end. A pass key stated back + correctly needs the right blocks selected at every depth. + + ``index_share`` makes the draft loop reuse the selection of its first step; the + acceptance length then must not fall below the floor taken from the same prompts + with the draft layer selecting on every step, and the reuse op must have run. + + With ``breakable`` the prompts are prefilled in 2048-token chunks, each replayed from + the top bucket: the continuations past the first chunk read the KV, index and recurrent + state the earlier replays wrote, which the single-chunk gsm8k prompts never do. + """ + engine = dict( + self.ENGINE_BREAKABLE_CHUNKED if breakable else self.ENGINE, + kv_cache_config=getattr(self, kv_cache), + ) + self._long_context_passkey( + monkeypatch, + self.MTP3, + engine, + index_share, + _PASSKEY_LENGTHS, + key_name, + f"mtp3 long context {kv_cache} breakable={breakable}", + ) + + @skip_not_sm103 + def test_breakable_prefill_matches_eager(self, monkeypatch): + """Greedy continuations of the same prompts, with the breakable prefill graph and without. + + MTP with three drafts, a 3072-token chunk budget, and prompts of a natural document + cut at the edges of the buckets 128 and 1024 (one token under, at, one over), at the + top bucket 3072 and one past it, and a 5000-token prompt prefilled in two chunks. A + batch of word problems and a 5000-token prompt then prefills while the earlier ones + decode, so replayed steps carry generation rows beside the context rows. + + A prompt that fills its bucket runs the same kernels on the same shapes as the eager + step, so its tokens must be identical. A padded prompt changes the row count every + GEMM sees, which moves the logits by what a different batch size does (about 0.2 nats + at most here); it must still agree on its first four tokens. The same engine given + the same prompts twice must repeat itself exactly. Replays are counted, since a + replay that never happens passes everything above. + """ + monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") + replays = _count_prefill_graph_replays(monkeypatch) + tokenizer = AutoTokenizer.from_pretrained(self.MODEL_PATH) + document = tokenizer.encode(_passkey_prompts(tokenizer, (6000,))[0][0]) + lengths = (127, 128, 129, 1023, 1024, 1025, 3072, 3073, 5000) + filled = (128, 1024, 3072) + # Word problems, answered in full, and one prompt of five thousand tokens prefilled in + # chunks while they decode. A word problem's first token can be a near tie between two + # openings, so its answer is checked, not its tokens. + problems = [ + tokenizer.apply_chat_template( + [{"role": "user", "content": f"{question} {_WORD_PROBLEM_SUFFIX}"}], + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + for question, _ in _WORD_PROBLEMS + ] + greedy = SamplingParams(max_tokens=24, temperature=0) + answering = SamplingParams(max_tokens=384, temperature=0) + + def run(engine): + with LLM( + self.MODEL_PATH, + speculative_config=self.MTP3, + **engine, + **_require_target(monkeypatch, **_VALIDATE), + ) as llm: + single = [ + llm.generate({"prompt_token_ids": document[:n]}, greedy, use_tqdm=False) + .outputs[0] + .token_ids + for n in lengths + ] + again = [ + llm.generate({"prompt_token_ids": document[:n]}, greedy, use_tqdm=False) + .outputs[0] + .token_ids + for n in lengths + ] + # The first requests are decoding when the long prompt arrives. + futures = [] + for index, prompt in enumerate(problems): + futures.append(llm.generate_async(prompt, answering)) + time.sleep(0.3) + if index == 1: + long_future = llm.generate_async( + {"prompt_token_ids": document[:5000]}, greedy + ) + time.sleep(0.3) + answers = [future.result().outputs[0].text for future in futures] + long_tokens = long_future.result().outputs[0].token_ids + return single, again, answers, long_tokens + + single, again, answers, long_tokens = run(dict(self.ENGINE_BREAKABLE, max_num_tokens=3072)) + assert {128, 256, 1024, 2048, 3072} <= set(replays), replays + assert single == again, "the same engine repeated itself with a different answer" + + replays.clear() + eager_single, _, eager_answers, eager_long_tokens = run( + dict(self.ENGINE, max_num_tokens=3072) + ) + assert not replays, replays + + for n, got, want in zip(lengths, single, eager_single): + head = len(want) if n in filled else 4 + assert got[:head] == want[:head], (n, got, want) + for texts in (answers, eager_answers): + wrong = [ + (question, answer, text[-120:]) + for (question, answer), text in zip(_WORD_PROBLEMS, texts) + if f"answer is {answer}" not in text.replace("**", "") + ] + assert not wrong, f"wrong or missing answers: {wrong}" + assert long_tokens[:4] == eager_long_tokens[:4], (long_tokens, eager_long_tokens) + + @skip_not_sm103 + def test_mtp7_long_context_passkey(self, monkeypatch): + """Pass key retrieval over 3k-12k-token prompts with seven drafts and IndexShare on. + + Seven drafts run the sparse path with eight rows per request in the verify step and + draft decode steps that reuse the first step's selection up to six positions past it, + where the groups completed since the capture displace two of its ids. 12k tokens past + the 8k default needs ``max_seq_len`` 16384. + """ + engine = dict(self.ENGINE, max_seq_len=16384) + self._long_context_passkey( + monkeypatch, + self.MTP7, + engine, + True, + _PASSKEY_LENGTHS_LONG, + "ACCEPTANCE_KEY_MTP7_LONG_CONTEXT", + "mtp7 long context", + ) + + @skip_not_sm103 + def test_mtp7_bf16_ssm_state_long_context_passkey(self, monkeypatch): + """``test_mtp7_long_context_passkey`` over a bf16 recurrent state. + + The state is rounded to bf16 at every cached-replay checkpoint commit, each time a + slot's history of 16 accepted tokens fills, and once per prefill chunk, so the rounding + adds up with the length of the prompt and of the answer; pass keys stated back over + 3k-12k tokens bound that drift. + """ + engine = dict(self.ENGINE, max_seq_len=16384, kv_cache_config=self.KV_CACHE_BF16_STATE) + self._long_context_passkey( + monkeypatch, + self.MTP7, + engine, + True, + _PASSKEY_LENGTHS_LONG, + "ACCEPTANCE_KEY_MTP7_BF16_STATE_LONG_CONTEXT", + "mtp7 bf16 state long context", + ) + + @pytest.mark.parametrize( + "forced, accepted_rows", + [ + pytest.param("1e-9", 1, id="all_rejected"), + pytest.param("7", 8, id="all_accepted"), + ], + ) + @skip_not_sm103 + def test_mtp7_forced_acceptance_extremes(self, monkeypatch, forced, accepted_rows): + """Seven drafts with every draft rejected, then with every draft accepted, at batch 32. + + ``TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS`` overrides the sampler's acceptance: + a vanishing fraction accepts the target's own token and nothing else, ``7`` accepts + the seven drafts as well. 32 requests over 2k-6k-token prompts run the sparse path, + the all-layer replay commit and the CUDA graph at 32 x 8 = 256 verify rows with the + most rows rolled back at every step (seven) and with none rolled back, a history that + then fills to its end every second step. The acceptance length, read from the + requests' own metrics, must be exactly the forced row count: the acceptance + bookkeeping of the verify step moved by the number of rows forced, not by what the + drafts say. + + The text under forced acceptance is not compared with anything: the drafts a forced + step accepts are not the target's tokens, and a run's tokens depend on the verify + step's row count through the GEMM and attention kernel selection, so a run with + another draft length or none differs from this one even with every draft rejected. + The rolled-back state itself is pinned bitwise by ``test_modeling_v2_spec_rollback``. + """ + tokenizer = AutoTokenizer.from_pretrained(self.MODEL_PATH) + prompts = [prompt for prompt, _ in _passkey_prompts(tokenizer, _FORCED_LENGTHS)][:32] + env = {"TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS": forced} + + with LLM( + self.MODEL_PATH, + speculative_config=self.MTP7, + **self.ENGINE, + **_require_target(monkeypatch, **env, **_VALIDATE), + ) as llm: + outputs = llm.generate( + prompts, + SamplingParams( + max_tokens=_FORCED_TOKENS, + temperature=0, + ignore_eos=True, + return_perf_metrics=True, + ), + use_tqdm=False, + ) + + assert [len(o.outputs[0].token_ids) for o in outputs] == [_FORCED_TOKENS] * len(prompts) + acceptance_length = _acceptance_length_from_requests(outputs, self.MTP7.max_draft_len) + print(f"[AL] mtp7 forced acceptance of {accepted_rows} rows: {acceptance_length:.3f}") + assert acceptance_length == pytest.approx(accepted_rows, abs=0.05) + + @skip_not_sm103 + def test_mtp1_gdn_legacy_verify(self, monkeypatch): + """MTP with one draft and ``TRTLLM_USE_GDN_REPLAY=0``: the per-step recurrent buffers. + + The engine runs the cached-replay verify step unless the variable is 0; this is the + only case that binds and runs the other one. A few word problems keep it short: their + answers say the recurrent state survives rejected drafts, and the acceptance length + says the drafts still land. + """ + monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") + tokenizer = AutoTokenizer.from_pretrained(self.MODEL_PATH) + prompts = [ + tokenizer.apply_chat_template( + [{"role": "user", "content": f"{question} {_WORD_PROBLEM_SUFFIX}"}], + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + for question, _ in _WORD_PROBLEMS + ] with LLM( self.MODEL_PATH, speculative_config=self.MTP1, - enable_iter_perf_stats=True, **self.ENGINE, - **_require_target(monkeypatch, **_GDN_LEGACY, **_VALIDATE), + **_require_target(monkeypatch, TRTLLM_USE_GDN_REPLAY="0", **_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) + outputs = llm.generate( + prompts, + SamplingParams(max_tokens=384, temperature=0, return_perf_metrics=True), + use_tqdm=False, + ) + bound = _target_core(llm)._bound + + assert bound.replay is None and bound.int_ssm, ( + "TRTLLM_USE_GDN_REPLAY=0 but the target bound the cached-replay verify state" + ) + wrong = [ + (question, answer, output.outputs[0].text[-120:]) + for (question, answer), output in zip(_WORD_PROBLEMS, outputs) + if f"answer is {answer}" not in output.outputs[0].text.replace("**", "") + ] + assert not wrong, f"wrong or missing answers: {wrong}" + acceptance_length = _acceptance_length_from_requests(outputs, self.MTP1.max_draft_len) + print(f"[AL] mtp1 gdn legacy acceptance_length = {acceptance_length:.3f}") + assert_acceptance_length(self.ACCEPTANCE_KEY_GDN_LEGACY, 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 + Block reuse is off here (on in the text 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", - ) + self._mmmu(monkeypatch, self.ENGINE) + + @skip_not_sm103 + def test_mtp1_mmmu_breakable_prefill(self, monkeypatch): + """``test_mtp1_mmmu`` over replayed prefill steps, against the same row. + + The vision rows are spliced into the embedding inside the first break of every + replay; a replay that kept the previous step's rows or dropped them would answer + from the text alone, which this score detects. + """ + self._mmmu(monkeypatch, self.ENGINE_BREAKABLE) + + def _mmmu(self, monkeypatch, engine): + engine = dict(engine, kv_cache_config=self.KV_CACHE_MULTIMODAL) # MMMU adds <|endoftext|> to the stop token, as the built-in model's # MMMU test does. @@ -180,7 +1065,7 @@ def test_mtp1_mmmu(self, monkeypatch): self.MODEL_PATH, speculative_config=self.MTP1, **engine, - **_require_target(monkeypatch, **_GDN_LEGACY, **_VALIDATE), + **_require_target(monkeypatch, **_VALIDATE), ) as llm: task = MMMU(self.MODEL_NAME) task.evaluate( @@ -188,3 +1073,56 @@ def test_mtp1_mmmu(self, monkeypatch): sampling_params=sampling_params, extra_evaluator_kwargs=dict(chat_template_kwargs=dict(enable_thinking=False)), ) + + @skip_not_sm103 + def test_mtp1_video(self, monkeypatch): + """Synthetic videos through the engine: the only case that sends a video. + + Three videos of a moving square (frames built in memory, since the container has no video + decoder) must each be answered with the right colour and direction. A request mixing an + image and a video, in either order, must reach the forward: its ``mm_item_order`` + manifest has to survive the input processor, the request and the admission check, which + refuses a mixed request without a usable one. A frame past the vision tower's patch + bound (a 3072x3072 image is 192x192 = 36864 patches) must fail as that request's own + error, and the request after it must still complete; ``max_seq_len`` is raised so that + image reaches the admission check rather than the sequence length check. The order of + the splice itself is pinned by ``test_modeling_v2_qwen4_exp_multimodal_embeds``. + """ + engine = dict(self.ENGINE, kv_cache_config=self.KV_CACHE_MULTIMODAL, max_seq_len=16384) + tokenizer = AutoTokenizer.from_pretrained(self.MODEL_PATH) + processor = AutoProcessor.from_pretrained( + self.MODEL_PATH, use_fast=True, trust_remote_code=True + ) + greedy = SamplingParams(max_tokens=96, temperature=0) + + def ask(text, items): + request = _media_request(tokenizer, processor, text, items) + return llm.generate_async(request, sampling_params=greedy).result().outputs[0].text + + with LLM( + self.MODEL_PATH, + speculative_config=self.MTP1, + **engine, + **_require_target(monkeypatch, **_VALIDATE), + ) as llm: + for name, (_, _, colour, directions) in _VIDEOS.items(): + answer = ask(_VIDEO_QUESTION, [("video", _synthetic_video(name))]).lower() + assert colour in answer and any(d in answer for d in directions), (name, answer) + + # Each label carries its own item's colour. + image = _synthetic_image((250, 200, 0)) + for items, video_colour in ( + ([("video", _synthetic_video("blue_leftward")), ("image", image)], "blue"), + ([("image", image), ("video", _synthetic_video("red_rightward"))], "red"), + ): + answer = ask(_MIXED_QUESTION, items).lower() + assert re.search(r"image:\W*yellow", answer), answer + assert re.search(rf"video:\W*{video_colour}", answer), answer + + with pytest.raises(RequestError, match="frames of 192x192 = 36864 patches"): + ask( + "What shape and colour is in the image?", + [("image", _synthetic_image(_RED, 3072))], + ) + answer = ask(_VIDEO_QUESTION, [("video", _synthetic_video("red_rightward"))]).lower() + assert "red" in answer and "right" in answer, answer diff --git a/tests/integration/defs/accuracy/references/acceptance_length.yaml b/tests/integration/defs/accuracy/references/acceptance_length.yaml index 45a750e9f47f..aaa9f6e21d62 100644 --- a/tests/integration/defs/accuracy/references/acceptance_length.yaml +++ b/tests/integration/defs/accuracy/references/acceptance_length.yaml @@ -123,21 +123,25 @@ ModelingV2DeepseekR10528Nvfp4Sm10xDefault::mtp3: TestGLM53FlashFP8::test_mtp: ref_al: 3.384 min_al: 3.215 -# Read by TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1:: -# test_mtp1_gsm8k_and_acceptance. Keyed by target and variant for the same -# reason as the ModelingV2DeepseekR10528Nvfp4Sm10xDefault::mtp3 entry above. +# Read by the acceptance-length cases of TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1 +# (test_qwen4_exp.py). Keyed by target and variant for the same reason as the +# ModelingV2DeepseekR10528Nvfp4Sm10xDefault::mtp3 entry above, and one key per variant: +# TRTLLM_POPULATE_ACCEPTANCE_LENGTH=1 on one case must not overwrite another's floor. # -# max_draft_len is 1 on this checkpoint (`mtp_num_hidden_layers: 1`, one module -# replayed once per step), so acceptance length lives in [1.0, 2.0]: 1.0 is a -# draft path that is rejected every time -- the exact failure a miscomputed MTP -# layer produces, with the text still correct and every accuracy gate green -- -# and 2.0 is every draft accepted. +# The checkpoint ships one MTP module (`mtp_num_hidden_layers: 1`), replayed max_draft_len +# times per step, so acceptance length lives in [1.0, 1 + max_draft_len]: 1.0 is a draft +# path that is rejected every time -- the exact failure a miscomputed MTP layer produces, +# with the text still correct and every accuracy gate green -- and 1 + max_draft_len is +# every draft accepted. Each min_al below is a hand-set collapse tripwire unless its +# comment says it is 95% of ref_al. + +# test_mtp1_gsm8k_and_acceptance ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp1: # Smoke-scale, and labelled as such: measured on GB300 over the boot prompt # pair, STOCK 1.7403 against modeling_v2's 1.8063 on the same prompts. The # number below is the STOCK reading, so the gate is not self-referential. - # There is no full-gsm8k reading yet: the accuracy runs were made without - # enable_iter_perf_stats, so they recorded no specDecodingStats. + # There is no full-gsm8k STOCK reading yet; modeling_v2's own over full gsm8k + # (1319 samples) is 1.971. ref_al: 1.7403 # Hand-set as a collapse tripwire, NOT the automatic 95% of ref_al -- 1.65 # would be a floor derived from a two-prompt smoke applied to a 1319-sample @@ -145,12 +149,121 @@ ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp1: # the variant's whole range: a collapsed or miscomputed draft path scores 1.0, # and anything at or above 1.5 is accepting at least half its drafts. # - # This unit has a disclosed deviation the gate is deliberately sized around: - # the draft layer runs dense attention rather than QSA, because - # `QSAAttentionMetadata` has no `prepare_for_draft_forward`. That costs - # acceptance on long prompts and costs nothing on short ones, so a floor set - # tight against a short-prompt reading would fail on the benchmark for a - # reason already written down. Replace both numbers with a full-gsm8k STOCK - # reading (TRTLLM_POPULATE_ACCEPTANCE_LENGTH=1 on the stock leg) when one is + # Replace both numbers with a full-gsm8k STOCK reading + # (TRTLLM_POPULATE_ACCEPTANCE_LENGTH=1 on the stock leg) when one is # measured, and tighten min_al then. min_al: 1.5 + +# test_mtp1_fp8_kv_gsm8k_and_acceptance: the same draft over an fp8 KV cache, with a floor of +# its own. Storage in e4m3 can cost drafts the bf16 cache keeps, and the bf16 floor of 1.5 +# would let a quarter of the acceptance go unnoticed. +ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp1_fp8_kv: + # modeling_v2's own full-gsm8k reading on GB300 (TRTLLM MoE, 1319 samples); the bf16 KV + # reading of the same run is 1.971, so fp8 storage costs 0.4%. + ref_al: 1.964 + # 95% of ref_al. + min_al: 1.866 + +# test_mtp1_nvfp4_kv_gsm8k_and_acceptance: the same draft over an NVFP4 KV cache, with a floor of +# its own for the same reason as the fp8 one. +ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp1_nvfp4_kv: + # modeling_v2's own full-gsm8k reading on GB300 (TRTLLM MoE, 1319 samples); the bf16 KV + # reading of the same run is 1.971, so NVFP4 storage costs 0.6%. + ref_al: 1.959 + # 95% of ref_al. + min_al: 1.861 + +# test_mtp3_gsm8k_and_acceptance: three drafts, bf16 KV, trtllm-gen experts. +ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3: + # The built-in implementation's full-gsm8k reading on GB300, so the gate is not + # self-referential. modeling_v2's own is 3.704 (MTP2: 2.876). + ref_al: 3.567 + # Hand-set as a collapse tripwire, NOT the automatic 95%: 3.0 is well below both readings + # (16% under the built-in's, 19% under the target's) yet above what a broken later draft + # position scores. The first draft alone is the one-draft value, ~1.97; the first two + # without the third score ~2.8; a draft path rejected every time scores 1.0. Each of those + # leaves the text correct and every accuracy gate green. + min_al: 3.0 + +# test_mtp3_cutedsl_gsm8k_and_acceptance: three drafts with the trunk's experts on CuteDSL. +ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_cutedsl: + # modeling_v2's own full-gsm8k reading; there is no built-in reading for this combination + # (the built-in model has no CuteDSL trunk), so the floor is the mtp3 tripwire. + ref_al: 3.690 + min_al: 3.0 + +# test_mtp3_long_context_passkey[index_share]: three drafts with IndexShare on, bf16 KV, over +# 13 prompts of 3k-6k tokens (the sparse attention path). Most of the output is a copy of the +# document's section headings, which the draft predicts almost always, so the readings sit +# at the range's ceiling of 4.0. +ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_long_context: + # The same prompts with the draft layer selecting on every step (IndexShare off): 4.000, + # equal to the IndexShare reading. + ref_al: 4.000 + # Hand-set, 2.5% below ref_al rather than the automatic 95% (3.8): a draft loop whose + # reused selection is deliberately wrong (the oldest 512 groups of every row) scores + # 3.760, a 6% drop the automatic floor would pass. + min_al: 3.9 + +# test_mtp3_long_context_passkey[fp8_kv]: the same prompts over an fp8 KV cache, IndexShare off. +ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_fp8_kv_long_context: + # 3.981 here; 3.973 with IndexShare on. + ref_al: 3.981 + # 2.5% below ref_al, as for mtp3_long_context. + min_al: 3.88 + +# test_mtp3_long_context_passkey[nvfp4_kv]: the same prompts over an NVFP4 KV cache, IndexShare off. +ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_nvfp4_kv_long_context: + # 3.971 here, against 4.000 over bf16 and 3.981 over fp8. + ref_al: 3.971 + # 2.5% below ref_al, as for mtp3_long_context. + min_al: 3.87 + +# test_mtp1_gdn_legacy_verify: one draft with TRTLLM_USE_GDN_REPLAY=0, over six word problems. +ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp1_gdn_legacy: + # modeling_v2's own reading over those problems, 384 new tokens each; the cached-replay + # verify path reads the same over full gsm8k (1.971). + ref_al: 1.968 + # 95% of ref_al. + min_al: 1.87 + +# test_mtp5_gsm8k_and_acceptance: five drafts, bf16 KV, trtllm-gen experts. +ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp5: + # modeling_v2's own full-gsm8k reading on GB300 (1319 samples; four drafts read 4.428, six + # 5.609); there is no built-in reading at this draft length. + ref_al: 5.063 + # 95% of ref_al: above the four-draft reading, so a fifth draft position that is never + # accepted fails it. + min_al: 4.810 + +# test_mtp7_gsm8k_and_acceptance: seven drafts, the longest the target runs. +ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp7: + # modeling_v2's own full-gsm8k reading on GB300, as for mtp5. + ref_al: 6.052 + # 95% of ref_al: above the six-draft reading (5.609), so a seventh draft position that is + # never accepted fails it. + min_al: 5.749 + +# test_mtp7_long_context_passkey: seven drafts with IndexShare on over 12 prompts of 3k-12k +# tokens (the sparse attention path), whose output mostly copies the document's section +# headings: the readings sit near the range's ceiling of 8. +ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp7_long_context: + ref_al: 7.774 + # 2.5% below ref_al, as for mtp3_long_context. + min_al: 7.580 + +# test_mtp3_bf16_ssm_state_gsm8k_and_acceptance: three drafts over a bf16 recurrent state. +ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_bf16_state: + # modeling_v2's own full-gsm8k reading on GB300; the fp32-state reading is + # 3.704, so the bf16 state costs nothing measurable. + ref_al: 3.699 + # 95% of ref_al, above the two-draft reading (~2.8): a third draft position that is never + # accepted fails it. + min_al: 3.514 + +# test_mtp7_bf16_ssm_state_long_context_passkey: the mtp7_long_context prompts over a bf16 +# recurrent state (7.774 over fp32). +ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp7_bf16_state_long_context: + ref_al: 7.857 + # 2.5% below ref_al, as for mtp7_long_context. + min_al: 7.661 diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index 4413b6b0f811..a2593fb2a263 100644 --- a/tests/integration/defs/accuracy/references/gsm8k.yaml +++ b/tests/integration/defs/accuracy/references/gsm8k.yaml @@ -299,6 +299,18 @@ Qwen/Qwen3.8-Flash-Next: - quant_algo: MIXED_PRECISION spec_dec_algo: MTP accuracy: 95.57 + # FP8 KV cache over the checkpoint's bf16 KV, measured on the built-in implementation: + # TP1 on GB300, MTP draft length 1, TRTLLM MoE, fp32 Mamba SSM cache. + - quant_algo: MIXED_PRECISION + kv_cache_quant_algo: FP8 + spec_dec_algo: MTP + accuracy: 95.94 + # NVFP4 KV cache over the checkpoint's bf16 KV, measured on the built-in implementation: + # TP1 on GB300, MTP draft length 1, TRTLLM MoE, fp32 Mamba SSM cache. + - quant_algo: MIXED_PRECISION + kv_cache_quant_algo: NVFP4 + spec_dec_algo: MTP + accuracy: 95.26 moonshotai/Kimi-K2.5: - quant_algo: NVFP4 kv_cache_quant_algo: FP8 diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index fd4554680089..d73cb66fc3a9 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -135,15 +135,50 @@ l0_b300: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_chunked_prefill[quant_dtype=fp8-kv_cache_reuse=True-fp8kv=True-overlap_scheduler=True] # Cover nvbugs 6084445 - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer[enable_chunked_prefill=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer[enable_chunked_prefill=True] - # TIMEOUT measured: 348 cases in 5m34s on one GB300, so 30 is 5x headroom. + # TIMEOUT measured: 2259 cases in 11m55s on one GB300 (the cpu_only files run on the CPU + # stage instead), so 30 is 2.5x headroom. - unittest/_torch/modeling_v2 --ignore=unittest/_torch/modeling_v2/comm TIMEOUT (30) - accuracy/modeling_v2/test_gpt_oss.py::TestModelingV2GptOss120bSm10xDefault::test_gsm8k[tp1] - # qwen4_exp target: gsm8k; gsm8k with MTP plus the acceptance-length floor, the - # only gate that sees a miscomputed draft layer; mmmu with MTP, the only gate - # that reaches the vision tower. + # qwen4_exp target: gsm8k, and gsm8k with the trunk's routed experts on CuteDSL; + # gsm8k with MTP plus the acceptance-length floor, the only gate that sees a + # miscomputed draft layer, at one draft over bf16 and fp8 KV, at three drafts + # alone and over CuteDSL experts, and at five and seven drafts (the verify step's + # six and eight rows per request); 3k-6k-token pass key prompts, the only cases that + # run the sparse attention path (IndexShare over bf16 KV, fp8 KV), and 3k-12k-token ones + # at seven drafts; seven drafts all rejected and all accepted by forced acceptance at + # batch 32 over long prompts; the recurrent-state verify path that + # TRTLLM_USE_GDN_REPLAY=0 selects; the recurrent state stored in bf16 without speculation, + # at three drafts and at seven drafts over long prompts; mmmu with MTP and synthetic + # videos, the only gates that reach the vision tower; and the opt-in breakable prefill + # CUDA graph: gsm8k, gsm8k with three drafts and the length floor, 3k-6k-token pass keys + # prefilled in replayed 2048-token chunks, mmmu, and replayed against eager greedy + # continuations at the bucket edges and with context and generation rows in one batch. - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_gsm8k + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_gsm8k_cutedsl_moe - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp1_gsm8k_and_acceptance + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp1_fp8_kv_gsm8k_and_acceptance + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp1_nvfp4_kv_gsm8k_and_acceptance + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_gsm8k_and_acceptance + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_cutedsl_gsm8k_and_acceptance + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp5_gsm8k_and_acceptance + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp7_gsm8k_and_acceptance + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_long_context_passkey[index_share] + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_long_context_passkey[fp8_kv] + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_long_context_passkey[nvfp4_kv] + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp7_long_context_passkey + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp7_forced_acceptance_extremes[all_rejected] + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp7_forced_acceptance_extremes[all_accepted] + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_gsm8k_bf16_ssm_state + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_bf16_ssm_state_gsm8k_and_acceptance + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp7_bf16_ssm_state_long_context_passkey + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp1_gdn_legacy_verify - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp1_mmmu + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp1_video + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_gsm8k_breakable_prefill_cuda_graph + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_breakable_prefill_gsm8k_and_acceptance + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_long_context_passkey[breakable_prefill] + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp1_mmmu_breakable_prefill + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_breakable_prefill_matches_eager - condition: ranges: system_gpu_count: diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 726469ed58ac..ff2b78f2a921 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -68,13 +68,21 @@ l0_cpu: - unittest/_torch/modeling # modeling_v2 routing and routing-table claims: config shapes and a monkeypatched # compute capability, no device. The default targets' shard tests: small CPU - # tensors and the CPU block_scale_interleave. The rest of - # unittest/_torch/modeling_v2 needs a GPU. + # tensors and the CPU block_scale_interleave. The qwen4_exp entries are the target's + # tests that run on the source text, stubbed weights and meta tensors: its multimodal + # splice and admission check on a stub tower, the KV pool contract, and its KV storage + # derivation and attention-call arguments. + # The rest of unittest/_torch/modeling_v2 needs a GPU. - unittest/_torch/modeling_v2/test_modeling_v2_claims.py - unittest/_torch/modeling_v2/test_modeling_v2_routing.py - unittest/_torch/modeling_v2/test_modeling_v2_deepseek_default_shards.py - unittest/_torch/modeling_v2/test_modeling_v2_gpt_oss_default_shards.py - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_dense_backend_library_order + - unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_multimodal_embeds.py + - unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_validate_request.py + - unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_pool_contract.py + - unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_kv_quant.py + - unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_thop_attention_calls.py - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_fp8_indexer_rejects_different_qk_norm_epsilons - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_moe_reduces_only_local_terms - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_decoder_layer_sets_post_fusion_from_moe_scheduler From 439b505612283e1de58f1e2058379a2cfe45321c Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Sun, 18 Oct 2026 22:59:58 -0700 Subject: [PATCH 04/28] [TRTLLM-16460][perf] Shared kernels for the Qwen3.8-Flash-Next target: low-M bf16 GEMM, FlashInfer GDN entries, PDL QSA index scores Catalog entries and QSA kernel changes the target's fast paths call, with their unit tests: - catalog/gemm/low_m_bf16_mm: decode-shaped bf16 projection, one kernel a call, its source picked per (M, N, K). - catalog/ssm/flashinfer_chunk_gated_delta_rule: FlashInfer's Blackwell chunked gated delta rule, the GDN prefill core. catalog/ssm/gated_delta_rule_decode_pretranspose: FlashInfer's single-token GDN decode over an fp32 slot pool, gates fused in. - catalog/attention/fused_qk_norm_rope: a cell for the target's geometry. - QSA kernels (shared with the in-tree QSA backend): the index-score kernel takes programmatic dependent launch (`use_pdl`) and can read its pages and keys before the wait (`prefetch`); both default off. The launch heuristics change for every caller: context-row tiles are 128 dot columns wide (was 16), batches past 4 rows use 64-block tiles (was 8 below 257 rows), and a program of packed context rows walks 8 tiles (was 32). Folds these commits of the review branch (their changes to the files of this commit): - qwen3.8 target: carry the decode-fusion perf work onto the feature target (62151b727b) Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../catalog/attention/fused_qk_norm_rope.py | 15 + .../modeling_v2/catalog/gemm/low_m_bf16_mm.py | 245 +++++++++++++++ .../modeling_v2/catalog/ssm/__init__.py | 3 + .../ssm/flashinfer_chunk_gated_delta_rule.py | 278 ++++++++++++++++++ .../gated_delta_rule_decode_pretranspose.py | 252 ++++++++++++++++ .../attention/backends/sparse/qsa/kernels.py | 196 +++++++++--- .../attention/sparse/qsa/test_qsa_sparse.py | 6 +- .../gemm/test_modeling_v2_low_m_bf16_mm.py | 155 ++++++++++ ...ng_v2_flashinfer_chunk_gated_delta_rule.py | 179 +++++++++++ ...v2_gated_delta_rule_decode_pretranspose.py | 167 +++++++++++ 10 files changed, 1460 insertions(+), 36 deletions(-) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/low_m_bf16_mm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/flashinfer_chunk_gated_delta_rule.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/gated_delta_rule_decode_pretranspose.py create mode 100644 tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_low_m_bf16_mm.py create mode 100644 tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py create mode 100644 tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_gated_delta_rule_decode_pretranspose.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/fused_qk_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/fused_qk_norm_rope.py index 73c6f700c652..e6dbc41e29ba 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/fused_qk_norm_rope.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/fused_qk_norm_rope.py @@ -103,6 +103,21 @@ class FusedQkNormRope(OpWrapper): "GEMMA_PREFILL_ATOL", spec=dict(tokens=8192, gemma_prefill_band=True, **_QWEN_MAIN), ), + # The main rope run in place on qwen4_exp's fused input projection, whose + # row is [q | k | v | gate | index | pad] = 55 heads of 256: everything + # after k is passed as v heads, which only sets the row stride. + Cell( + why="qwen4_exp fused-row main rope, 1 token: 29 v heads only set the row stride", + spec=dict(tokens=1, **{**_QWEN_MAIN, "num_heads_v": 29}), + ), + Cell( + why="qwen4_exp fused-row main rope, 128 tokens: decode at max concurrency", + spec=dict(tokens=128, **{**_QWEN_MAIN, "num_heads_v": 29}), + ), + Cell( + why="qwen4_exp fused-row main rope, 8192 tokens: prefill, under GEMMA_PREFILL_ATOL", + spec=dict(tokens=8192, gemma_prefill_band=True, **{**_QWEN_MAIN, "num_heads_v": 29}), + ), Cell( why="qwen4_exp index rope (mRoPE), 1 token: one k and one v head, the " "narrowest layout; 5 warps of an 8-warp block", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/low_m_bf16_mm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/low_m_bf16_mm.py new file mode 100644 index 000000000000..72fda36b187e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/low_m_bf16_mm.py @@ -0,0 +1,245 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Decode-shaped bf16 projection: one kernel per call, its source picked per (M, N, K). + +At decode M a dense bf16 projection streams its weight once, and cuBLAS +reaches only 0.55-0.65 of HBM bandwidth on 0.25-63 MB weights (and ~4.5 us +flat on the small ones). FlashInfer's Blackwell low-M CuTe-DSL kernels -- the +register-prefetch "direct" kernel and the tcgen05 split-K kernel, the two +`flashinfer.mm_bf16(backend="cute-dsl")` dispatches to and +`tensorrt_llm/_torch/modules/low_m_gemm.py` races -- are faster at M <= 32 on +most such shapes. Which one, with which tactic, is a per-(N, K) table measured +per M bucket on GB300 (cold L2, CUDA-graph replay); every shape the table does +not list, and every M above it, stays on the trtllm cuBLASLt op. The rows +today are qwen3.8-flash-next-nvfp4's projections (hidden 2560, tp1). + +FlashInfer API: ``flashinfer.gemm.kernels.dense_bf16_gemm_direct.run_direct_dense`` and +``flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk.run_splitk_dense``. +""" + +from __future__ import annotations + +import functools +from typing import NamedTuple + +import torch +from flashinfer.gemm.kernels import dense_bf16_gemm_direct as _fi_direct +from flashinfer.gemm.kernels import dense_bf16_gemm_sm100_splitk as _fi_splitk + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* + +from .._op import Arch, Cell, OpWrapper, true_fp32_matmul + + +class _Plan(NamedTuple): + """`kind` is "direct", "splitk" or "cublas"; `tactic` is the kernel's own tactic tuple.""" + + kind: str + tactic: tuple = () + + +def _d(block_size: int, outputs_per_block: int) -> tuple: + """A direct-kernel choice; rows_per_block is the kernel's default for the runtime M.""" + return ("direct", block_size, outputs_per_block) + + +def _s(mma_m: int, mma_n: int, split_k: int, ab_stages: int) -> tuple: + return ("splitk", mma_m, mma_n, split_k, ab_stages) + + +#: (N, K) -> ((M upper bound, choice), ...) in ascending M. An M above the last +#: bound, and any (N, K) not listed, runs on cuBLAS. Every row below is a +#: qwen3.8-flash-next-nvfp4 tp1 projection; a target with other shapes adds its +#: own rows after measuring them. Each choice beat cublas_mm by >= 5% at the +#: bound's M in +#: qwen3.8-flash-next-nvfp4-perf-opt/05-iterations/candidates/low-m-gemm/microbench/sweep.json +#: or, for the attention in-projection concat and the PLE value projection, +#: .../candidates/lowm-integrate/microbench/sweep_new.json +#: (GB300, cold L2, graph replay, pdl off). M=3 uses the M=4 choice, and so on. +_TABLE: dict[tuple[int, int], tuple[tuple[int, tuple], ...]] = { + # GDN in_qkv: 16 q + 16 k heads x 128, 48 v heads x 128 + (10240, 2560): ((1, _d(64, 1)), (2, _d(64, 2)), (8, _s(128, 8, 1, 6))), + # GDN in_z + (6144, 2560): ((1, _d(64, 1)), (2, _d(64, 2))), + # GDN in_a / in_b (48 v heads) + (48, 2560): ((32, _d(64, 1)),), + # GDN out and attention o: [6144 -> 2560] + (2560, 6144): ( + (2, _d(256, 2)), + (8, _s(64, 8, 3, 10)), + (16, _s(64, 16, 3, 8)), + (32, _s(64, 32, 3, 8)), + ), + # attention q (q and its output gate interleaved per head) + (12288, 2560): ((2, _d(64, 2)), (4, _s(128, 8, 1, 6)), (8, _s(128, 16, 1, 6))), + # attention k / v (2 kv heads x 256) and the MoE router (512 experts) + (512, 2560): ((2, _d(64, 2)), (8, _d(64, 1)), (16, _d(32, 1)), (32, _s(64, 8, 4, 10))), + # QSA index projection: 4 index heads x 128 + 1 key x 128 + (640, 2560): ((4, _d(64, 1)), (8, _d(32, 1)), (16, _s(64, 8, 4, 10)), (32, _s(64, 16, 4, 8))), + # shared expert gate_up, and the same with its gate row appended (padded to 1296) + (1280, 2560): ((2, _d(64, 1)), (4, _d(64, 2)), (8, _s(64, 8, 4, 10))), + (1296, 2560): ((2, _d(64, 1)), (4, _d(64, 2)), (8, _s(64, 8, 4, 10))), + # shared expert gate [2560 -> 1] + (1, 2560): ((32, _d(64, 1)),), + # the four GDN input projections concatenated (10240 + 6144 + 48 + 48) + (16480, 2560): ((1, _d(64, 2)),), + # the attention input projections concatenated: q|gate (12288), k, v (512 each), + # index (640), zero-padded to 55 heads of 256 (14080) + (14080, 2560): ((2, _d(64, 2)), (8, _s(128, 8, 1, 6))), + # PLE value projection (decoder layer 1 only); the PLE key projection is [10240, 2560] + (2560, 2560): ((1, _d(64, 1)), (2, _d(64, 2)), (4, _d(32, 2)), (8, _s(64, 8, 2, 10))), +} +# Measured and left on cuBLAS: (2560, 640) shared-expert down (the direct kernel +# cannot tile K=640; split-K and TGV were within 4% of cuBLAS), the GDN concat +# (16480, 2560) at M >= 2 (direct 3.6-5.3% ahead at M=2 across two sweeps, behind +# from M=4), the attention concat (14080, 2560) at M >= 16, and M >= 64 +# everywhere (outside both kernels' M <= 32 policy). Not servable at all: the +# hyper-connection mixer's first projection (fp32 output; both kernels write +# bf16 only) and its up projection [10240, 320] (K=320 tiles neither kernel). + + +@functools.lru_cache(maxsize=None) +def plan(m: int, n: int, k: int) -> _Plan: + """The kernel and tactic `LowMBf16Mm` launches for an [M, K] x [K, N] product.""" + for bound, choice in _TABLE.get((n, k), ()): + if m <= bound: + if choice[0] == "direct": + t = _fi_direct.DirectTactic( + choice[1], choice[2], _fi_direct.default_tactic(m, n, k).rows_per_block + ) + _fi_direct.validate_tactic(t, m, n, k) + return _Plan("direct", tuple(choice[1:]) + (t.rows_per_block,)) + _fi_splitk.validate_tactic(_fi_splitk.SplitKTactic(*choice[1:]), m, n, k) + return _Plan("splitk", tuple(choice[1:])) + return _Plan("cublas") + + +def _cell(why: str, m: int, n: int, k: int, kind: str, pdl: bool = False) -> Cell: + """`kind` is the branch `plan` must pick for this shape; the test holds it to that.""" + return Cell(why=why, spec=dict(m=m, n=n, k=k, kind=kind, pdl=pdl)) + + +class LowMBf16Mm(OpWrapper): + """`mat_a @ mat_b` in bf16 with fp32 accumulation: one kernel launch per call. + + A drop-in for `catalog/gemm/cublas_mm` without bias or `out_dtype`: `mat_a` + is dense row-major `[M, K]`, `mat_b` dense column-major `[K, N]` (a stored + `[N, K]` weight's `.t()`), the result a fresh row-major `[M, N]` bf16 tensor. + `plan(M, N, K)` picks the kernel on the host from three sources: + + - `flashinfer.gemm.kernels.dense_bf16_gemm_direct.run_direct_dense` -- SIMT, + the weight row's whole K prefetched into registers; wins at M <= 2 and on + narrow N; + - `flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk.run_splitk_dense` -- + tcgen05, swap-AB, split-K reduced across a cluster in shared memory (no + extra launch); wins at M 4-32 on the wide shapes; + - `torch.ops.trtllm.cublas_mm` -- everything else, including every M > 32. + + `pdl` launches the FlashInfer kernels with programmatic dependent launch: + they wait on the previous grid (`griddepcontrol.wait`) before reading + `mat_a`, so it is safe after any kernel; it changes nothing on the cuBLAS + branch. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # Every cell but the one off the table is a qwen3.8-flash-next-nvfp4 tp1 + # projection shape; `kind` pins the branch `plan` takes for it. + CELLS: tuple[Cell, ...] = ( + _cell("qwen3.8 decode M=1 GDN in_qkv: direct", 1, 10240, 2560, "direct"), + _cell("qwen3.8 decode M=2 GDN in_z: direct, 2-row tile", 2, 6144, 2560, "direct"), + _cell("qwen3.8 decode M=1 GDN in_a/in_b, N=48: direct", 1, 48, 2560, "direct"), + _cell("qwen3.8 decode M=32 GDN in_a/in_b: direct, 8-row tiles x4", 32, 48, 2560, "direct"), + _cell("qwen3.8 decode M=1 o_proj K=6144: direct, 256-thread K", 1, 2560, 6144, "direct"), + _cell("qwen3.8 decode M=8 o_proj: split-K 3 over the cluster", 8, 2560, 6144, "splitk"), + _cell("qwen3.8 decode M=32 o_proj: split-K, mma_n 32", 32, 2560, 6144, "splitk"), + _cell("qwen3.8 decode M=3 attention q: split-K, M off-bucket", 3, 12288, 2560, "splitk"), + _cell("qwen3.8 decode M=8 attention q: split-K, mma_n 16", 8, 12288, 2560, "splitk"), + _cell("qwen3.8 decode M=16 router / k / v N=512: direct", 16, 512, 2560, "direct"), + _cell("qwen3.8 decode M=24 router / k / v: split-K 4", 24, 512, 2560, "splitk"), + _cell("qwen3.8 decode M=16 index proj N=640: split-K", 16, 640, 2560, "splitk"), + _cell("qwen3.8 decode M=4 shared gate_up: direct", 4, 1280, 2560, "direct"), + _cell("qwen3.8 decode M=8 shared gate_up + gate row: split-K", 8, 1296, 2560, "splitk"), + _cell("qwen3.8 decode M=1 shared gate N=1: direct", 1, 1, 2560, "direct"), + _cell("qwen3.8 decode M=1 GDN in-proj concat N=16480: direct", 1, 16480, 2560, "direct"), + _cell("qwen3.8 decode M=2 GDN in-proj concat: cuBLAS branch", 2, 16480, 2560, "cublas"), + _cell( + "qwen3.8 decode M=2 attention in-proj concat N=14080: direct", 2, 14080, 2560, "direct" + ), + _cell("qwen3.8 decode M=3 attention in-proj concat: split-K", 3, 14080, 2560, "splitk"), + _cell("qwen3.8 decode M=8 attention in-proj concat: split-K", 8, 14080, 2560, "splitk"), + _cell("qwen3.8 decode M=16 attention in-proj concat: cuBLAS", 16, 14080, 2560, "cublas"), + _cell("qwen3.8 decode M=1 PLE value N=K=2560: direct", 1, 2560, 2560, "direct"), + _cell("qwen3.8 decode M=4 PLE value: direct, 32-thread K", 4, 2560, 2560, "direct"), + _cell("qwen3.8 decode M=8 PLE value: split-K 2", 8, 2560, 2560, "splitk"), + _cell("qwen3.8 decode M=1 shared down K=640: cuBLAS branch", 1, 2560, 640, "cublas"), + _cell("qwen3.8 decode M=16 attention q: cuBLAS branch", 16, 12288, 2560, "cublas"), + _cell("qwen3.8 decode M=128: cuBLAS branch above M=32", 128, 10240, 2560, "cublas"), + _cell("qwen3.8 o_proj prefill-shaped M=8192: cuBLAS branch", 8192, 2560, 6144, "cublas"), + _cell("decode M=1, an (N, K) not in the table: cuBLAS", 1, 1024, 2048, "cublas"), + _cell("qwen3.8 pdl on, direct branch", 1, 10240, 2560, "direct", pdl=True), + _cell("qwen3.8 pdl on, split-K branch", 8, 2560, 6144, "splitk", pdl=True), + ) + + note = """ + Metadata consumed: none. Preconditions: bf16 operands, `mat_a` dense + row-major, `mat_b` dense column-major, both 32-byte aligned (a freshly + allocated activation and a `.t()` weight view are); the FlashInfer kernels + raise ValueError on a layout or alignment they cannot read, the cuBLAS + branch does not (see `is_valid`). + + Compilation: the FlashInfer kernels are CuTe-DSL, JIT-compiled in-process + on the first call per (direct: M, N, K, tactic, pdl; split-K: tactic, pdl) + and cached for the process. That first call must happen before CUDA-graph + capture -- the engine's warmup forward at each captured batch size does it. + + Host cost: the FlashInfer branches spend tens of microseconds of Python per + call (argument validation, DLPack conversion), which a CUDA graph removes + and an eager step does not. Route only graph-captured decode calls here. + + Numerics: every branch accumulates in fp32 and rounds once; the branches + differ from each other only in summation order. + + Hard limits + M FlashInfer branches only at M <= 32 (both kernels' policy) boundary 32 driven at 32 + (cells) and 128, 8192 (cuBLAS cells) + N*K largest FlashInfer-branch weight 16480*2560 = 4.2e7 elements boundary 2^31 not probed (48x below it) + """ + + def __call__(self, mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool = False) -> torch.Tensor: + m, k = mat_a.shape + n = mat_b.shape[1] + p = plan(m, n, k) + if p.kind == "direct": + out = torch.empty((m, n), dtype=mat_a.dtype, device=mat_a.device) + return _fi_direct.run_direct_dense( + mat_a, mat_b, out, pdl, _fi_direct.DirectTactic(*p.tactic) + ) + if p.kind == "splitk": + out = torch.empty((m, n), dtype=mat_a.dtype, device=mat_a.device) + return _fi_splitk.run_splitk_dense( + mat_a, mat_b, None, out, pdl, _fi_splitk.SplitKTactic(*p.tactic) + ) + return torch.ops.trtllm.cublas_mm(mat_a, mat_b, None, None, 0, None) + + def reference( + self, mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool = False + ) -> torch.Tensor: + """fp32-accumulated, TF32 off; `pdl` is a launch attribute and is ignored.""" + with true_fp32_matmul(): + return (mat_a.float() @ mat_b.float()).to(mat_a.dtype) + + def is_valid(self, mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool = False) -> None: + # The cuBLAS branch ignores both leading dimensions: a row-padded + # mat_a or a column-padded mat_b (a slice of a wider tensor) is read as + # if dense -- silently a different product. The FlashInfer branches + # raise on those. A dense row-major mat_b is outside the contract too + # (cuBLAS and the direct kernel raise on it; split-K computes it). + assert mat_a.dtype == torch.bfloat16 and mat_b.dtype == torch.bfloat16, "bf16 operands only" + assert mat_a.is_contiguous(), "mat_a must be dense row-major [M, K]" + assert mat_b.stride(0) == 1 and mat_b.stride(1) == mat_b.shape[0], ( + "mat_b must be dense column-major [K, N] (e.g. weight.t())" + ) + + +low_m_bf16_mm = LowMBf16Mm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/__init__.py new file mode 100644 index 000000000000..bbf1df41d716 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""SSM entries.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/flashinfer_chunk_gated_delta_rule.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/flashinfer_chunk_gated_delta_rule.py new file mode 100644 index 000000000000..7504fb20757a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/flashinfer_chunk_gated_delta_rule.py @@ -0,0 +1,278 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Chunked gated delta rule on FlashInfer's Blackwell kernel: the GDN prefill core. + +Wraps the in-tree module callable +``tensorrt_llm._torch.modules.fla.flashinfer_chunk.chunk_gated_delta_rule`` +(tier-2 sourcing), the adapter in front of the installed FlashInfer's +``flashinfer.chunk_gated_delta_rule`` (CuTe-DSL tcgen05 kernel on sm_100/103). +One call is: ``g -> exp(g)``, a Triton gather of the indexed pool slots into a +packed fp32 ``[N, HV, V, K]`` buffer, the FlashInfer kernel, and a Triton +scatter of the final states back into the same slots. Call-compatible with the +vendored Triton ``chunk_gated_delta_rule`` (``modules/fla/chunk``). + +FlashInfer API: ``flashinfer.chunk_gated_delta_rule`` (``flashinfer.gdn_prefill``), +reached through the adapter above. +""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch + +from tensorrt_llm._torch.modules.fla.flashinfer_chunk import ( + chunk_gated_delta_rule as _fi_chunk_gated_delta_rule, +) + +from .._op import Arch, Cell, OpWrapper + +#: qwen3.8-flash-next-nvfp4's Gated DeltaNet at tp1: (H, HV, K, V) -- 16 q/k +#: heads, 48 v heads, head dims 128. +QWEN38_GEOMETRY = (16, 48, 128, 128) +#: Largest slot count the cache manager was seen to allocate the recurrent pool +#: at for qwen3.8-flash-next-nvfp4 (10 .. 461 over twelve runtime-value dumps). +ENGINE_SLOTS = 461 + + +def _engine_slots(num_seqs: int) -> tuple[int, ...]: + """`num_seqs` distinct slots of the engine pool, slot 0 and the last live.""" + return ((0, ENGINE_SLOTS - 1) + tuple(i * 13 + 4 for i in range(num_seqs)))[:num_seqs] + + +class FlashinferChunkGatedDeltaRule(OpWrapper): + """Run the gated delta rule over whole sequences, returning `(o, final_state)`. + + Per token `t` and value head, `h` the `[V, K]` state, `q_h`/`k_h` the key + heads broadcast (value head `i` reads key head `i // (HV // H)`), fp32 + accumulation: + + h *= exp(g[t]); u = (v[t] - h k_h[t]) * beta[t]; h += u k_h[t]^T + o[t] = h (scale * q_h[t]) + + computed 64 tokens at a time on tensor cores (WY form, bf16 operands) -- + the same recurrence as the Triton op, a different sum. `g` is in log space + (the adapter exponentiates it); `scale=None` is `K ** -0.5`; + `use_qk_l2norm_in_kernel` normalizes q/k first (a Triton `l2norm_fwd`). + + One packed row (`q.shape[0] == 1`) split at `cu_seqlens` (required). + `initial_state` is required: a pool `[S, HV, V, K]`; sequence `n` starts + from `initial_state[initial_state_indices[n]]`. With + `inplace_indexed_state_update=True` its final state is written back into + that slot and `final_state` is `None`; otherwise `output_final_state` + returns a fresh `[N, HV, V, K]` in the pool's dtype and the pool is not + written. + + Fusion boundary: gate exponentiation, optional q/k norm, scale, slot + gather, recurrence, output, slot scatter. The conv, q/k/v/g/beta + preparation, slot allocation and the gated norm on `o` are the caller's. + + Returns `(o, final_state)`: `o` `[1, T, HV, V]` in `q.dtype` (`output` + itself when given). + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: geometry (H, HV, K, V); seqlens; `slots`-row pool with `indices` + # (or _engine_slots); inplace (True); output_final_state; output (None | + # "buffer"); l2norm (False) with normalize_qk (True: caller-normalized). + # Every cell but the last is qwen3.8-flash-next-nvfp4's geometry. + CELLS: tuple[Cell, ...] = tuple( + Cell( + why=f"qwen3.8 prefill, 461-slot pool, seqlens {list(s)}", + spec=dict( + geometry=QWEN38_GEOMETRY, + seqlens=s, + slots=ENGINE_SLOTS, + indices=_engine_slots(len(s)), + ), + ) + for s in ((1,), (64,), (200, 37, 91), (1024,), (8192,), (2048, 2048, 2048, 2048)) + ) + ( + Cell( + why="qwen3.8 context chunk + 1-token rows [8, 1, 1, 1], 461-slot pool", + spec=dict( + geometry=QWEN38_GEOMETRY, + seqlens=(8, 1, 1, 1), + slots=ENGINE_SLOTS, + indices=_engine_slots(4), + ), + ), + Cell( + why="qwen3.8 output= buffer [1, T, HV, V], 6-slot pool", + spec=dict( + geometry=QWEN38_GEOMETRY, + seqlens=(128, 64), + slots=6, + indices=(4, 1), + output="buffer", + ), + ), + Cell( + why="qwen3.8 no in-place: final states returned, pool untouched", + spec=dict( + geometry=QWEN38_GEOMETRY, + seqlens=(96, 48), + slots=5, + indices=(3, 0), + inplace=False, + output_final_state=True, + ), + ), + Cell( + why="qwen3.8 use_qk_l2norm_in_kernel on raw q/k", + spec=dict( + geometry=QWEN38_GEOMETRY, + seqlens=(100,), + slots=3, + indices=(2,), + l2norm=True, + normalize_qk=False, + ), + ), + Cell( + why="HV/H = 2 (16 q/k heads, 32 v heads), seqlens [200, 37]", + spec=dict(geometry=(16, 32, 128, 128), seqlens=(200, 37), slots=4, indices=(3, 0)), + ), + ) + + note = """ + Metadata consumed: none. + + Call site as qwen3.8-flash-next-nvfp4's `_gdn_block` makes it (in place of + the Triton `chunk_gated_delta_rule`): + the same positional arguments, `(q, k, v, g, beta, None, ssm_states, + si, True, False, qsl_long, False, False, None)` with `head_first=False` + in the 12th slot, after `fused_gdn_post_conv` (q/k already normalized, + contiguous `[1, T, H, 128]`; `g` fp32 log-space; `beta` fp32). Give it + context rows only: single-token rows are slow here. Measured per layer + at 8065 + 127 x 1 tokens: 1078.8 us with every row in this op, 399.4 us + with the 127 decode rows sent to `gated_delta_rule_decode_pretranspose` + instead (both halves plus the concatenation). + + Memory per call: a packed fp32 copy of the named slots in and out (2 x + 3 MiB per sequence), `exp(g)`, and the output. q/k/v/beta are not copied + when already contiguous / fp32. + + JIT and streams: the first call per process compiles the CuTe-DSL + kernel in-process (seconds), and the installed FlashInfer keeps one TMA + workspace per device, so two streams must not run this op concurrently. + Needs CUDA >= 13 torch (FlashInfer raises otherwise). + + Numerics (measured vs an fp64 recurrence, 1 x 8192 tokens, qwen3.8 + geometry): o rel-L2 4.06e-3, max abs 1.26e-3; state rel-L2 3.31e-3 -- + the Triton op's are + 4.00e-3 / 1.13e-3 / 3.26e-3. Not bitwise to it. No cold-call defect: + nothing in the call is autotuned. + + Loud, from the callable: `head_first=True`, `q.shape[0] != 1`, H_q != H_k, + fp32 q, non-fp32 `g`, `cu_seqlens` or `initial_state` None, in-place + without indices. + + Not checked: indices distinct and in `[0, S)`; q/k unit-norm (by the + caller or the flag). + + Not certified: `state_workspace=`, bf16/fp16 pool, `scale` other than the + default, fp16 operands, K or V other than 128. + + Hard limits + none probed. + """ + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: Optional[float] = None, + initial_state: Optional[torch.Tensor] = None, + initial_state_indices: Optional[torch.Tensor] = None, + inplace_indexed_state_update: bool = False, + output_final_state: bool = False, + cu_seqlens: Optional[torch.Tensor] = None, + head_first: bool = False, + use_qk_l2norm_in_kernel: bool = False, + output: Optional[torch.Tensor] = None, + state_workspace: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + return _fi_chunk_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + initial_state_indices=initial_state_indices, + inplace_indexed_state_update=inplace_indexed_state_update, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + head_first=head_first, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + output=output, + state_workspace=state_workspace, + ) + + def reference( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: Optional[float] = None, + initial_state: Optional[torch.Tensor] = None, + initial_state_indices: Optional[torch.Tensor] = None, + inplace_indexed_state_update: bool = False, + output_final_state: bool = False, + cu_seqlens: Optional[torch.Tensor] = None, + head_first: bool = False, + use_qk_l2norm_in_kernel: bool = False, + output: Optional[torch.Tensor] = None, + state_workspace: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + ) -> tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor]: + """`(o, final_state, initial_state as the call leaves it)`. + + fp64 token-at-a-time recurrence, rounded once: `o` to `q.dtype`, the + states to the pool's dtype. `final_state` is `None` on the in-place + path, as the callable returns it. Mutates nothing; `output`, + `head_first` and `state_workspace` are ignored (they choose where + results land).""" + dt = torch.float64 + _, _, num_k_heads, head_k = q.shape + num_v_heads = v.shape[2] + repeats = num_v_heads // num_k_heads + scale = head_k**-0.5 if scale is None else scale + bounds = cu_seqlens.tolist() + slots = initial_state_indices.tolist() + out = torch.zeros(1, q.shape[1], num_v_heads, v.shape[3], dtype=dt, device=q.device) + finals = [] + for n in range(len(bounds) - 1): + state = initial_state[slots[n]].to(dt) + for t in range(bounds[n], bounds[n + 1]): + q_t, k_t = q[0, t].to(dt), k[0, t].to(dt) + if use_qk_l2norm_in_kernel: + q_t = q_t / torch.sqrt(q_t.pow(2).sum(-1, keepdim=True) + 1e-6) + k_t = k_t / torch.sqrt(k_t.pow(2).sum(-1, keepdim=True) + 1e-6) + q_t = q_t.repeat_interleave(repeats, dim=0) * scale + k_t = k_t.repeat_interleave(repeats, dim=0) + state = state * g[0, t].to(dt).exp()[:, None, None] + u = (v[0, t].to(dt) - (state * k_t[:, None, :]).sum(-1)) * beta[0, t].to(dt)[ + :, None + ] + state = state + u[:, :, None] * k_t[:, None, :] + out[0, t] = (state * q_t[:, None, :]).sum(-1) + finals.append(state) + pool = initial_state.clone() + final_state = None + if inplace_indexed_state_update: + for slot, state in zip(slots, finals): + pool[slot] = state.to(pool.dtype) + elif output_final_state: + final_state = torch.stack(finals).to(initial_state.dtype) + return out.to(q.dtype), final_state, pool + + +flashinfer_chunk_gated_delta_rule = FlashinferChunkGatedDeltaRule() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/gated_delta_rule_decode_pretranspose.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/gated_delta_rule_decode_pretranspose.py new file mode 100644 index 000000000000..20a11f7cf5db --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/gated_delta_rule_decode_pretranspose.py @@ -0,0 +1,252 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Gated DeltaNet single-token decode over an fp32 slot pool, gates fused in. + +Wraps FlashInfer's public +``flashinfer.gdn_decode.gated_delta_rule_decode_pretranspose`` (tier-3 +sourcing, the installed FlashInfer). With an fp32 pool and ``T == 1`` it +runs the CuTe-DSL "pretranspose" kernel (``run_pretranspose_decode``): eight +CTAs per (sequence, v-head) state, each streaming 16 of its 128 V rows through +shared memory, the pool read and written in place through +``initial_state_indices``. A bf16 pool or ``T > 1`` route elsewhere inside the +same callable and are not certified here. + +FlashInfer API: ``flashinfer.gdn_decode.gated_delta_rule_decode_pretranspose``. +""" + +from __future__ import annotations + +from typing import Optional + +import flashinfer.gdn_decode +import torch + +from .._op import Arch, Cell, OpWrapper + +#: qwen3.8-flash-next-nvfp4's Gated DeltaNet at tp1: (H, HV, K, V) -- 16 q/k +#: heads, 48 v heads, head dims 128. +QWEN38_GEOMETRY = (16, 48, 128, 128) +#: Largest slot count the cache manager was seen to allocate the recurrent pool +#: at for qwen3.8-flash-next-nvfp4 (10 .. 461 over twelve runtime-value dumps). +ENGINE_SLOTS = 461 + + +def _scattered(n: int, slots: int = ENGINE_SLOTS) -> tuple[int, ...]: + """`n` distinct slots of a `slots`-row pool, both ends first, then a + stride-37 walk over the interior (37 is coprime to every pool size here).""" + rest = tuple((i * 37 + 4) % (slots - 2) + 1 for i in range(max(n - 2, 0))) + return ((slots - 1, 0) + rest)[:n] + + +class GatedDeltaRuleDecodePretranspose(OpWrapper): + """One gated-delta-rule step per sequence, gates computed in-kernel. + + Row `n` of `q`/`k` `[B, 1, H, K]`, `v` `[B, 1, HV, V]`, `a`/`b` `[B, 1, HV]`; + v-head `hv` reads q/k head `hv // (HV // H)`; state `h` `[V, K]` (K + innermost) read from `initial_state[initial_state_indices[n]]`, in fp32: + + g = -exp(A_log) * softplus(a + dt_bias) (softplus beta 1, threshold 20) + beta = sigmoid(b) + q, k = x * rsqrt(sum(x^2) + 1e-6) (only with use_qk_l2norm) + h *= exp(g); h += k (beta * (v - h k))^T; o = h (scale * q) + + `scale=None` is `K ** -0.5`. The final `h` goes back to the same slot, or + to `output_state_indices[n]` when given. A negative index skips the row: + no state read or write, `o` written as zero. + + Fusion boundary: both gates, the optional q/k norm, scale, recurrence, + indexed gather and scatter. The conv, the q/k/v split and the gated RMS + norm on `o` are the caller's. + + Returns `(output, initial_state)` as the callable does: `output` + `[B, 1, HV, V]` (the `output` buffer itself when given, bf16; else a fresh + tensor in `q.dtype`) and the pool, advanced in place. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: geometry (H, HV, K, V), m (rows), slots (pool rows), indices (None + # -> _scattered(m, slots)), qkv ("views": column views of one + # [m, 2*H*K + HV*V] conv output, as a decode forward splits it | + # "contiguous"), output ("buffer" | None), pool ("contiguous" | "padded": + # slot stride HV*V*K + 64 | "layer_view": pool[1] of a [3, slots, ...] + # cache). Every cell but the last is qwen3.8-flash-next-nvfp4's geometry. + CELLS: tuple[Cell, ...] = tuple( + Cell( + why=f"qwen3.8 decode M={m}, 461-slot pool, scattered slots", + spec=dict(geometry=QWEN38_GEOMETRY, m=m, slots=ENGINE_SLOTS), + ) + for m in (1, 2, 8, 32, 64, 128) + ) + ( + Cell( + why="qwen3.8 output None: fresh q.dtype output, M=4", + spec=dict(geometry=QWEN38_GEOMETRY, m=4, slots=ENGINE_SLOTS, output=None), + ), + Cell( + why="qwen3.8 CUDA-graph pad rows (-1) between slot 0 and 460", + spec=dict( + geometry=QWEN38_GEOMETRY, + m=4, + slots=ENGINE_SLOTS, + indices=(0, -1, ENGINE_SLOTS - 1, -1), + ), + ), + Cell( + why="qwen3.8 padded slot stride (HV*V*K + 64), M=8", + spec=dict(geometry=QWEN38_GEOMETRY, m=8, slots=40, pool="padded"), + ), + Cell( + why="qwen3.8 pool[1] of a [3, slots, HV, V, K] cache, M=4", + spec=dict(geometry=QWEN38_GEOMETRY, m=4, slots=24, pool="layer_view"), + ), + Cell( + why="qwen3.8 contiguous q/k/v instead of conv-output views, M=16", + spec=dict(geometry=QWEN38_GEOMETRY, m=16, slots=ENGINE_SLOTS, qkv="contiguous"), + ), + Cell( + why="HV/H = 2 (16 q/k heads, 32 v heads), M=8, 40-slot pool", + spec=dict(geometry=(16, 32, 128, 128), m=8, slots=40), + ), + ) + + note = """ + Metadata consumed: none. + + Call site as qwen3.8-flash-next-nvfp4's decode `_gdn_block` makes it (in + place of `fused_sigmoid_gating_delta_rule_update`): `q=reshape(qv, [M, 1, 16, 128])`, + `k=...`, `v=reshape(vv, [M, 1, 48, 128])` -- column views of the conv + output, read through their strides -- `state=None`, `A_log`/`dt_bias` + fp32 `[48]`, `a`/`b` the `[M, 48]` projections viewed `[M, 1, 48]`, + `scale=None`, `output=` a bf16 `[M, 1, 48, 128]` buffer, `use_qk_l2norm= + True`, `initial_state=` the layer's fp32 pool `[slots, 48, 128, 128]`, + `initial_state_indices=` int32 slots. `output.view(1, M, 48, 128)` is what + `_gdn_epilogue` reads. + + Preconditions not checked by anything: + - fp32 pool, K-innermost `(V, K)` trailing axes, slot stride a multiple of + 4 elements (asserted) -- a transposed `[S, HV, K, V]` pool at K == V has + the same shape and is accepted. + - Indices distinct and in `[0, S)` or negative. An index >= S is an + unchecked device access: measured, slot 460 of a 40-slot pool raised + `cudaErrorIllegalAddress` on the next sync and poisoned the process. + - `A_log` fp32 (asserted), `dt_bias` fp32 or bf16 (asserted). + + Loud: every tensor's base pointer must be 16-byte aligned (the kernel is + compiled with `assumed_align=16`). Measured: indices taken as + `idx[n_c:]` of an int32 vector at an offset not a multiple of 4 raise + `RuntimeError: Tensor data pointer is not aligned to 16 bytes.` -- so a + mixed step's decode half, `si_all` split after its context rows, needs + a copy. The decode-only step's `si_all` starts at offset 0. + + JIT and capture: the first call per (`T, H, HV, K, V, q.dtype, scale, + use_qk_l2norm`, pool inner strides) compiles a CuTe-DSL kernel in-process + (seconds); the first call per batch size allocates two small int32 aux + tensors, unused on the pool path. Both must happen in eager warmup, before + graph capture. Capture-safe after that: no host sync, no allocation when + `output` is given (one `torch.zeros` per call otherwise). + + Numerics (measured at qwen3.8's geometry, fp64 reference): o max abs 1.2e-4 + at M=128 (|o| <= 6.2e-2), equal to the in-tree Triton op's + (`modules/fla/fused_sigmoid_gating_recurrent`); written state max abs + 2.6e-7 (Triton 2.0e-7). Not bitwise to the Triton op: rsqrt(sum + eps) vs + 1 / (sqrt(sum) + eps) for the q/k norm, and fastmath exp/log. + + Not certified: bf16/fp16 pool (another kernel), `T > 1` (routes to + `gated_delta_rule_mtp`), `state=` direct path, `output_state_indices`, + fp16 operands, K or V other than 128. + + Hard limits + none probed. + """ + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + state: Optional[torch.Tensor], + A_log: torch.Tensor, + a: torch.Tensor, + dt_bias: torch.Tensor, + b: torch.Tensor, + scale: Optional[float] = None, + output: Optional[torch.Tensor] = None, + use_qk_l2norm: bool = True, + initial_state: Optional[torch.Tensor] = None, + initial_state_indices: Optional[torch.Tensor] = None, + output_state_indices: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + return flashinfer.gdn_decode.gated_delta_rule_decode_pretranspose( + q=q, + k=k, + v=v, + state=state, + A_log=A_log, + a=a, + dt_bias=dt_bias, + b=b, + scale=scale, + output=output, + use_qk_l2norm=use_qk_l2norm, + initial_state=initial_state, + initial_state_indices=initial_state_indices, + output_state_indices=output_state_indices, + ) + + def reference( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + state: Optional[torch.Tensor], + A_log: torch.Tensor, + a: torch.Tensor, + dt_bias: torch.Tensor, + b: torch.Tensor, + scale: Optional[float] = None, + output: Optional[torch.Tensor] = None, + use_qk_l2norm: bool = True, + initial_state: Optional[torch.Tensor] = None, + initial_state_indices: Optional[torch.Tensor] = None, + output_state_indices: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """`(o, pool_after)`: fp64, rounded once -- `o` to `output`'s dtype (bf16 + when given, else `q.dtype`), the state to the pool's. `state=` (no pool) + is treated as a pool of `B` rows read and written in order. Mutates + nothing.""" + dt = torch.float64 + batch, _, num_k_heads, k_dim = q.shape + hv = v.shape[2] + rep = hv // num_k_heads + scale = k_dim**-0.5 if scale is None else scale + pool_in = initial_state if initial_state is not None else state + if initial_state is not None: + reads = initial_state_indices.tolist() + else: + reads = list(range(batch)) + writes = reads if output_state_indices is None else output_state_indices.tolist() + pool = pool_in.clone() + o = torch.zeros(batch, 1, hv, v.shape[3], dtype=dt, device=q.device) + neg_a = -torch.exp(A_log.to(dt)) + for n in range(batch): + if reads[n] < 0: + continue + q_n = q[n, 0].to(dt).repeat_interleave(rep, dim=0) + k_n = k[n, 0].to(dt).repeat_interleave(rep, dim=0) + if use_qk_l2norm: + q_n = q_n * torch.rsqrt(q_n.pow(2).sum(-1, keepdim=True) + 1e-6) + k_n = k_n * torch.rsqrt(k_n.pow(2).sum(-1, keepdim=True) + 1e-6) + x = a[n, 0].to(dt) + dt_bias.to(dt) + softplus = torch.where(x <= 20.0, torch.log1p(torch.exp(x)), x) + beta = torch.sigmoid(b[n, 0].to(dt)) + h = pool_in[reads[n]].to(dt) * torch.exp(neg_a * softplus)[:, None, None] + u = (v[n, 0].to(dt) - torch.einsum("hvk,hk->hv", h, k_n)) * beta[:, None] + h = h + u[:, :, None] * k_n[:, None, :] + o[n, 0] = torch.einsum("hvk,hk->hv", h, q_n * scale) + if writes[n] >= 0: + pool[writes[n]] = h.to(pool.dtype) + out_dtype = output.dtype if output is not None else q.dtype + return o.to(out_dtype), pool + + +gated_delta_rule_decode_pretranspose = GatedDeltaRuleDecodePretranspose() diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/kernels.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/kernels.py index dec4259363a8..358b5f424b43 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/kernels.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/kernels.py @@ -1119,7 +1119,24 @@ def _qsa_paged_index_scores_kernel( BLOCK_N: tl.constexpr, TILES_PER_PROGRAM: tl.constexpr, ONLY_VISIBLE_BLOCKS: tl.constexpr, + USE_PDL: tl.constexpr = False, + PREFETCH: tl.constexpr = False, ): + if PREFETCH: + # One tile a program: its pages and keys are read before the wait (see + # ``triton_qsa_paged_index_scores``), q after it. + tl.static_assert(USE_PDL and TILES_PER_PROGRAM == 1 and ONLY_VISIBLE_BLOCKS) + _qsa_paged_index_scores_prefetch_tile( + q, index_cache, block_table, query_positions, request_indices, output, score_scale, + q_stride_row, q_stride_head, q_stride_dim, cache_stride_page, cache_stride_token, + cache_stride_dim, block_table_stride_request, block_table_stride_page, + output_stride_row, output_stride_block, NUM_INDEX_HEADS, INDEX_HEAD_DIM, BLOCK_D, + TOKENS_PER_BLOCK, COMPRESS_RATIO, MAX_COMPRESSED_BLOCKS, BLOCK_H, BLOCK_N, + ) # fmt: skip + return + if USE_PDL: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() row = tl.program_id(0) request = tl.load(request_indices + row).to(tl.int64) query_position = tl.load(query_positions + row) @@ -1188,6 +1205,95 @@ def _qsa_paged_index_scores_kernel( ) +@triton.jit +def _qsa_paged_index_scores_prefetch_tile( + q, + index_cache, + block_table, + query_positions, + request_indices, + output, + score_scale, + q_stride_row: tl.constexpr, + q_stride_head: tl.constexpr, + q_stride_dim: tl.constexpr, + cache_stride_page: tl.constexpr, + cache_stride_token: tl.constexpr, + cache_stride_dim: tl.constexpr, + block_table_stride_request: tl.constexpr, + block_table_stride_page: tl.constexpr, + output_stride_row: tl.constexpr, + output_stride_block: tl.constexpr, + NUM_INDEX_HEADS: tl.constexpr, + INDEX_HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, + TOKENS_PER_BLOCK: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + MAX_COMPRESSED_BLOCKS: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_N: tl.constexpr, +): + """One generation row's tile of visible compressed blocks under PDL, its static + operands read before ``gdc_wait``: the row's metadata, the tile's pages and every + key but the one a decode step's own prologue may have just written -- the last + visible block's, compressed at its anchor when the row's position ends a group -- + which is read again after the wait, with q.""" + tl.extra.cuda.gdc_launch_dependents() + row = tl.program_id(0) + request = tl.load(request_indices + row).to(tl.int64) + query_position = tl.load(query_positions + row) + visible_blocks = (query_position + 1) // COMPRESS_RATIO + first_tile = tl.program_id(1) + if first_tile * BLOCK_N >= visible_blocks: + tl.extra.cuda.gdc_wait() + return + dim_offsets = tl.arange(0, BLOCK_D) + head_offsets = tl.arange(0, BLOCK_H) + block_columns = first_tile * BLOCK_N + tl.arange(0, BLOCK_N) + valid_blocks = (block_columns < MAX_COMPRESSED_BLOCKS) & (block_columns < visible_blocks) + anchor_positions = block_columns * COMPRESS_RATIO + COMPRESS_RATIO - 1 + logical_pages = anchor_positions // TOKENS_PER_BLOCK + token_in_page = anchor_positions % TOKENS_PER_BLOCK + physical_pages = tl.load( + block_table + + request * block_table_stride_request + + logical_pages * block_table_stride_page, + mask=valid_blocks, + other=0, + ).to(tl.int64) + key_ptrs = ( + index_cache + + physical_pages[:, None] * cache_stride_page + + token_in_page[:, None] * cache_stride_token + + dim_offsets[None, :] * cache_stride_dim + ) + key_mask = valid_blocks[:, None] & (dim_offsets[None, :] < INDEX_HEAD_DIM) + keys = tl.load(key_ptrs, mask=key_mask, other=0.0, eviction_policy="evict_first") + tl.extra.cuda.gdc_wait() + fresh = (block_columns == visible_blocks - 1) & ((query_position + 1) % COMPRESS_RATIO == 0) + keys = tl.where( + fresh[:, None], tl.load(key_ptrs, mask=key_mask & fresh[:, None], other=0.0), keys + ) + query_values = tl.load( + q + + row * q_stride_row + + dim_offsets[:, None] * q_stride_dim + + head_offsets[None, :] * q_stride_head, + mask=(dim_offsets[:, None] < INDEX_HEAD_DIM) & (head_offsets[None, :] < NUM_INDEX_HEADS), + other=0.0, + ) + per_head_scores = tl.dot(keys, query_values, out_dtype=tl.float32) + scores = tl.sum( + tl.where(head_offsets[None, :] < NUM_INDEX_HEADS, tl.maximum(per_head_scores, 0.0), 0.0), + axis=1, + ) + tl.store( + output + row * output_stride_row + block_columns * output_stride_block, + scores * score_scale, + mask=valid_blocks, + ) + + @triton.jit def _qsa_paged_index_scores_qtile_kernel( q, @@ -1248,6 +1354,11 @@ def _qsa_paged_index_scores_qtile_kernel( other=0.0, ) column_offsets = tl.arange(0, BLOCK_N) + # Each tile row's (lane's) visible bound, indexed by lane instead of by column. + lane_rows = row_base + tl.arange(0, QUERY_TILE) + lane_present = lane_rows < num_rows + lane_positions = tl.load(query_positions + lane_rows, mask=lane_present, other=0) + lane_visible = tl.where(lane_present, (lane_positions + 1) // COMPRESS_RATIO, 0) # A gather is only shareable across rows that page through the same block # table, so score one request at a time. Context rows arrive packed per @@ -1266,6 +1377,10 @@ def _qsa_paged_index_scores_qtile_kernel( request = tl.max(tl.where(dot_columns == leader, requests, 0)) group = pending & (requests == request) pending = pending & (requests != request) + # The lanes whose columns are in this group (none on a pass past the last group). + lane_group = ( + tl.max(tl.reshape(group.to(tl.int32), [QUERY_TILE, NUM_INDEX_HEADS]), axis=1) > 0 + ) group_visible = tl.max(tl.where(group, visible_blocks, 0)) last_tile = tl.minimum( first_tile + TILES_PER_PROGRAM, @@ -1300,37 +1415,34 @@ def _qsa_paged_index_scores_qtile_kernel( tl.dot(keys, query_values, out_dtype=tl.float32), 0.0, ) - for lane in tl.static_range(QUERY_TILE): - row_columns = group & (column_row == lane) - scored = tl.max(row_columns.to(tl.int32)) > 0 - row_visible = tl.max(tl.where(row_columns, visible_blocks, 0)) - scores = tl.sum( - tl.where(row_columns[None, :], positive_scores, 0.0), - axis=1, + # Each row's heads are adjacent dot columns: sum them per row. + scores = tl.sum( + tl.reshape(positive_scores, [BLOCK_N, QUERY_TILE, NUM_INDEX_HEADS]), axis=2 + ) + output_ptrs = ( + output + + lane_rows[None, :] * output_stride_row + + block_columns[:, None] * output_stride_block + ) + row_valid = block_columns[:, None] < lane_visible[None, :] + if ONLY_VISIBLE_BLOCKS: + tl.store( + output_ptrs, + scores * score_scale, + mask=lane_group[None, :] & row_valid, ) - output_ptrs = ( - output - + (row_base + lane) * output_stride_row - + block_columns * output_stride_block + else: + tl.store( + output_ptrs, + tl.where(row_valid, scores * score_scale, -float("inf")), + mask=lane_group[None, :] & (block_columns[:, None] < MAX_COMPRESSED_BLOCKS), ) - row_valid = block_columns < row_visible - if ONLY_VISIBLE_BLOCKS: - tl.store( - output_ptrs, - scores * score_scale, - mask=scored & row_valid, - ) - else: - tl.store( - output_ptrs, - tl.where(row_valid, scores * score_scale, -float("inf")), - mask=scored & (block_columns < MAX_COMPRESSED_BLOCKS), - ) - - -# ``tl.dot`` needs at least 16 columns on SM90+, so a narrower model dimension -# leaves the rest of every tensor-core tile computed and discarded. -_INDEX_SCORE_TILE_WIDTH = 16 + + +# Dot columns of a context-row index-scoring tile: whole query rows of +# ``num_index_heads`` columns each. A wide tile serves more rows with each +# gathered key tile (``tl.dot`` needs at least 16 columns on SM90+). +_INDEX_SCORE_TILE_WIDTH = 128 def _qsa_index_scores_query_tile(num_index_heads: int) -> int: @@ -1358,6 +1470,8 @@ def triton_qsa_paged_index_scores( compress_ratio: int, only_visible_blocks: bool = False, context_rows: bool = False, + use_pdl: bool = False, + prefetch: bool = False, ) -> torch.Tensor: """Score compressed keys directly in the paged side cache. @@ -1368,6 +1482,17 @@ def triton_qsa_paged_index_scores( ``context_rows`` declares that consecutive rows mostly share a request, so one gather of compressed keys can serve a tile of query rows. Generation batches hold one row per request and would only lose parallelism. + + ``use_pdl`` launches the generation-row kernel with programmatic dependent + launch: it waits for its predecessor before its first read and releases its + dependents at once, so every dependent must itself wait before reading the + scores. + + ``prefetch`` (with ``use_pdl``, ``only_visible_blocks`` and generation-sized + batches) reads the pages and compressed keys before the wait: sound only when the + predecessor writes no compressed key but, at a row whose position ends a group, the + last visible block's -- the decode step's own K/V-and-index prologue -- which is + read again after the wait. """ if q.ndim != 3: raise ValueError("QSA index queries must have shape [rows, heads, dim]") @@ -1401,10 +1526,12 @@ def triton_qsa_paged_index_scores( ) if rows == 0: return output - # Decode favors many narrow score tiles; large prefill batches reuse each - # query across a deeper column range to reduce launch and gather overhead. + # A decode batch of a few rows favors many narrow score tiles; past that, + # each program's fixed cost (the query load, the page lookups) is spread + # over a wider tile. Large prefill batches also reuse each query across a + # deeper column range to reduce launch and gather overhead. is_prefill_sized = rows > 256 - block_n = 64 if is_prefill_sized else 8 + block_n = 64 if rows > 4 else 8 tiles_per_program = 8 if is_prefill_sized else 1 query_tile = 1 if context_rows and is_prefill_sized: @@ -1412,7 +1539,7 @@ def triton_qsa_paged_index_scores( if query_tile > 1: # Row packing reduces the grid, so each program walks a deeper column # range to retain enough work per launch. - tiles_per_program = 32 + tiles_per_program = 8 grid = ( triton.cdiv(rows, query_tile), triton.cdiv(max_compressed_blocks, block_n * tiles_per_program), @@ -1488,7 +1615,10 @@ def triton_qsa_paged_index_scores( BLOCK_N=block_n, TILES_PER_PROGRAM=tiles_per_program, ONLY_VISIBLE_BLOCKS=only_visible_blocks, + USE_PDL=use_pdl, + PREFETCH=prefetch and use_pdl and tiles_per_program == 1 and only_visible_blocks, num_warps=2, + launch_pdl=use_pdl, ) return output diff --git a/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py index 6dc868d1c5fb..ddc36f14be14 100644 --- a/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py +++ b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py @@ -777,7 +777,7 @@ def test_qsa_paged_index_scores_match_reference(rows: int, head_dim: int) -> Non @pytest.mark.parametrize( "num_index_heads, expected", - [(1, 16), (2, 8), (4, 4), (8, 2), (16, 1), (3, 1), (12, 1), (24, 1), (0, 1)], + [(1, 128), (2, 64), (4, 32), (8, 16), (16, 8), (128, 1), (3, 1), (12, 1), (24, 1), (0, 1)], ) def test_qsa_index_scores_query_tile_divides_the_dot_width( num_index_heads: int, expected: int @@ -785,7 +785,7 @@ def test_qsa_index_scores_query_tile_divides_the_dot_width( query_tile = _qsa_index_scores_query_tile(num_index_heads) assert query_tile == expected if query_tile > 1: - assert query_tile * num_index_heads == 16 + assert query_tile * num_index_heads == 128 def _paged_index_scores_reference( @@ -824,7 +824,7 @@ def test_qsa_paged_index_scores_query_tile_matches_row_wise( head_dim = 16 # The query tile only replaces the row-wise kernel above the prefill row # boundary, so keep the launch there or the comparison is vacuous. - assert _qsa_index_scores_query_tile(num_heads) == 4 + assert _qsa_index_scores_query_tile(num_heads) == 32 if rows_per_request == "many": # Request boundaries deliberately land inside a query tile. lengths = [101, 157, 42] diff --git a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_low_m_bf16_mm.py b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_low_m_bf16_mm.py new file mode 100644 index 000000000000..c0b4dfd69ac4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_low_m_bf16_mm.py @@ -0,0 +1,155 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the low_m_bf16_mm catalog entry. + +Every cell names the branch `plan` must pick, so a table edit that moves a +shape to another kernel fails here before it can change a number silently. +Beyond the cells: CUDA-graph capture and replay after warmup (the way the +decode forward runs it), the wrong-variant discrimination of the band, and the +guards. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.low_m_bf16_mm import ( + low_m_bf16_mm as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.low_m_bf16_mm import plan + +__extra_import_path__ = [".."] +from _validating import ( # noqa: E402,F401 — path above; fixture by name + ULP, + assert_within_ulp, + certifying, + validating, +) + +assert torch.cuda.is_available(), "low_m_bf16_mm requires a CUDA device" + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Accumulation order only: 4 bf16 ulp element-wise against the row scale, + 2 ulp of relative RMS -- the band `cublas_mm` uses, for the same reason. + Every branch sums the same bf16 products in fp32 and rounds once; what is + left between kernel and fp32 reference is the final bf16 rounding (half an + ulp) plus the summation order.""" + assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0, ulp=ULP[torch.bfloat16]) + + +def _build(m, n, k, seed): + """`mat_b` is a `[N, K]` weight's `.t()`, as a target stores its projections.""" + g = torch.Generator(device="cuda").manual_seed(seed) + mat_a = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + weight = (torch.randn(n, k, generator=g, device="cuda") * k**-0.5).to(torch.bfloat16) + return mat_a, weight.t() + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 + certifying(op) + s = cell.spec + assert plan(s["m"], s["n"], s["k"]).kind == s["kind"], plan(s["m"], s["n"], s["k"]) + mat_a, mat_b = _build(s["m"], s["n"], s["k"], seed=zlib.crc32(cell.why.encode()) % 2**31) + out = op(mat_a, mat_b, s["pdl"]) + assert out.shape == (s["m"], s["n"]) and out.dtype == torch.bfloat16 and out.is_contiguous() + _compare(out, op.reference(mat_a, mat_b, s["pdl"])) + + +@pytest.mark.parametrize("pdl", [False, True]) +def test_graph_replay_after_warmup(pdl) -> None: + """Decode runs this under CUDA graphs: compile in warmup, capture, then replay + with new activations copied into the captured input. One call per branch, in + a chain, so a FlashInfer kernel follows cuBLAS and vice versa.""" + shapes = [(1, 10240, 2560), (8, 2560, 6144), (1, 2560, 640), (8, 12288, 2560)] + assert [plan(*s).kind for s in shapes] == ["direct", "splitk", "cublas", "splitk"] + ins = [_build(m, n, k, seed=i) for i, (m, n, k) in enumerate(shapes)] + static = [a.clone() for a, _ in ins] + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(2): + [op(a, b, pdl) for a, (_, b) in zip(static, ins)] + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outs = [op(a, b, pdl) for a, (_, b) in zip(static, ins)] + for trial in range(3): + g = torch.Generator(device="cuda").manual_seed(100 + trial) + for a in static: + a.copy_(torch.randn(a.shape, generator=g, device="cuda").to(a.dtype)) + graph.replay() + torch.cuda.synchronize() + for a, (_, b), out in zip(static, ins, outs): + _compare(out, op.reference(a, b)) + + +@pytest.mark.parametrize("m,n,k", [(1, 10240, 2560), (8, 2560, 6144)]) +def test_band_discriminates_wrong_variants(m, n, k) -> None: + """The band admits the kernel and rejects the near misses a wrong layout or + a dropped K-split would produce: half of K missing, a one-column shift of + the output, a 3% scale error.""" + mat_a, mat_b = _build(m, n, k, seed=7) + out = op(mat_a, mat_b) + ref = op.reference(mat_a, mat_b) + _compare(out, ref) + half = mat_a.clone() + half[:, k // 2 :] = 0 + for wrong in (op.reference(half, mat_b), torch.roll(ref, 1, dims=1), ref * 1.03): + with pytest.raises(AssertionError): + _compare(out, wrong) + + +def _padded(m, n, k): + """A row-padded mat_a and a column-padded mat_b: dense along K, leading + dimension 2K. Both are views a caller gets by slicing a wider tensor.""" + a = torch.randn(m, 2 * k, dtype=torch.bfloat16, device="cuda")[:, :k] + w = (torch.randn(n, 2 * k, device="cuda") * k**-0.5).to(torch.bfloat16)[:, :k] + return a, w.t() + + +def test_guard_refuses_a_padded_mat_b() -> None: + """Why the guard exists: the cuBLAS branch ignores mat_b's leading dimension + and silently computes a different product.""" + a, b = _padded(4, 640, 640) + a = a.contiguous() + assert plan(4, 640, 640).kind == "cublas" and b.stride() == (1, 1280) + with pytest.raises(AssertionError): + _compare(op(a, b), op.reference(a, b)) + with pytest.raises(AssertionError, match="column-major"): + with validating(op): + op(a, b) + + +def test_guard_refuses_a_padded_mat_a() -> None: + """Same for mat_a's row stride on the cuBLAS branch.""" + a, b = _padded(4, 640, 640) + b = b.t().contiguous().t() + assert plan(4, 640, 640).kind == "cublas" and a.stride() == (1280, 1) + with pytest.raises(AssertionError): + _compare(op(a, b), op.reference(a, b)) + with pytest.raises(AssertionError, match="row-major"): + with validating(op): + op(a, b) + + +def test_flashinfer_branches_reject_padded_operands_loudly() -> None: + """The FlashInfer branches validate layouts themselves and raise.""" + a, b = _padded(1, 10240, 2560) + assert plan(1, 10240, 2560).kind == "direct" + with pytest.raises(ValueError): + op(a.contiguous(), b) + a, b = _padded(8, 2560, 6144) + assert plan(8, 2560, 6144).kind == "splitk" + with pytest.raises(ValueError): + op(a, b.t().contiguous().t()) + + +def test_direct_branch_rejects_a_row_major_mat_b_loudly() -> None: + a = torch.randn(1, 2560, dtype=torch.bfloat16, device="cuda") + b = torch.randn(2560, 10240, dtype=torch.bfloat16, device="cuda") + assert plan(1, 10240, 2560).kind == "direct" + with pytest.raises(ValueError): + op(a, b) diff --git a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py b/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py new file mode 100644 index 000000000000..ca055aec966b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py @@ -0,0 +1,179 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the flashinfer_chunk_gated_delta_rule catalog entry (FlashInfer GDN prefill).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.flashinfer_chunk_gated_delta_rule import ( # noqa: E501 + QWEN38_GEOMETRY, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.flashinfer_chunk_gated_delta_rule import ( # noqa: E501 + flashinfer_chunk_gated_delta_rule as op, +) + +__extra_import_path__ = [".."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "flashinfer_chunk_gated_delta_rule requires a CUDA device" + +DEV = "cuda" +#: The band the Triton chunk op is certified at (rtol bf16 default; atol 2e-3 +#: on `o`, 2e-2 on states), derived for a chunked WY reformulation with bf16 +#: operands. This op is the same class of sum: measured against fp64 at +#: 1 x 8192 tokens, o max abs 1.26e-3 and state max abs 8.6e-3 (Triton op: +#: 1.13e-3 / 7.7e-3), inside it. +RTOL, ATOL_OUTPUT, ATOL_STATE = 1.6e-2, 2e-3, 2e-2 + + +def _compare(out: tuple, ref: tuple) -> None: + """`(o, final_state, pool)` against `reference`'s triple; `None` must match.""" + for got, want, atol in zip(out, ref, (ATOL_OUTPUT, ATOL_STATE, ATOL_STATE)): + assert (got is None) == (want is None), (got is None, want is None) + if got is None: + continue + assert got.dtype == want.dtype and got.shape == want.shape, ( + (got.dtype, got.shape), + (want.dtype, want.shape), + ) + assert torch.isfinite(got).all() + torch.testing.assert_close(got.float(), want.float(), rtol=RTOL, atol=atol) + + +def _build(spec, seed): + """Keyword arguments in the domain a GDN post-conv step produces: + unit-norm q/k (unless the cell feeds raw ones), g < 0 in log space, beta a sigmoid.""" + gen = torch.Generator(device=DEV).manual_seed(seed) + num_k_heads, num_v_heads, head_k_dim, head_v_dim = spec["geometry"] + seqlens = spec["seqlens"] + total = sum(seqlens) + q = torch.randn(1, total, num_k_heads, head_k_dim, generator=gen, device=DEV) + k = torch.randn(1, total, num_k_heads, head_k_dim, generator=gen, device=DEV) + if spec.get("normalize_qk", True): + q = q / torch.sqrt(q.pow(2).sum(-1, keepdim=True) + 1e-6) + k = k / torch.sqrt(k.pow(2).sum(-1, keepdim=True) + 1e-6) + q, k = q.bfloat16(), k.bfloat16() + v = torch.randn( + 1, total, num_v_heads, head_v_dim, generator=gen, device=DEV, dtype=torch.bfloat16 + ) + a_log = torch.log(torch.rand(num_v_heads, generator=gen, device=DEV) * 15 + 1) + x = torch.randn(1, total, num_v_heads, generator=gen, device=DEV) + torch.randn( + num_v_heads, generator=gen, device=DEV + ) + g = (-torch.exp(a_log) * torch.nn.functional.softplus(x)).float() + beta = torch.rand(1, total, num_v_heads, generator=gen, device=DEV).sigmoid().float() + pool = ( + torch.randn(spec["slots"], num_v_heads, head_v_dim, head_k_dim, generator=gen, device=DEV) + * 0.1 + ) + cu = torch.tensor([0] + torch.tensor(seqlens).cumsum(0).tolist(), dtype=torch.long, device=DEV) + indices = torch.tensor(spec["indices"], dtype=torch.int32, device=DEV) + output = None + if spec.get("output") == "buffer": + output = torch.full( + (1, total, num_v_heads, head_v_dim), -7.0, device=DEV, dtype=torch.bfloat16 + ) + return dict( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=None, + initial_state=pool, + initial_state_indices=indices, + inplace_indexed_state_update=spec.get("inplace", True), + output_final_state=spec.get("output_final_state", False), + cu_seqlens=cu, + head_first=False, + use_qk_l2norm_in_kernel=spec.get("l2norm", False), + output=output, + state_workspace=None, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) + pool = args["initial_state"] + pool_before = pool.clone() + ref = op.reference(**args) + + o, final_state = op(**args) + + _compare((o, final_state, pool), ref) + if args["output"] is not None: + assert o.data_ptr() == args["output"].data_ptr(), "output= must be written in place" + named = set(args["initial_state_indices"].tolist()) + unnamed = [s for s in range(pool.shape[0]) if s not in named] + assert torch.equal(pool[unnamed], pool_before[unnamed]), "unnamed slots written" + if not args["inplace_indexed_state_update"]: + assert torch.equal(pool, pool_before), "pool written without in-place update" + + +def test_repeated_calls_are_bitwise() -> None: + """Two identical calls from the same start give identical output and state.""" + results = [] + for _ in range(2): + args = _build( + dict(geometry=QWEN38_GEOMETRY, seqlens=(200, 37, 91), slots=6, indices=(5, 0, 2)), + seed=7, + ) + o, _ = op(**args) + results.append((o.clone(), args["initial_state"].clone())) + assert torch.equal(results[0][0], results[1][0]) and torch.equal(results[0][1], results[1][1]) + + +def test_matches_in_tree_triton_chunk_op() -> None: + """Same arguments through the vendored Triton chunk op this entry stands in + for: o and the written slots agree at the band. The Triton op's cold + indexed call is spent on scratch first.""" + from tensorrt_llm._torch.modules.fla.chunk import chunk_gated_delta_rule as triton_op + + def run_triton(a, pool): + return triton_op( + q=a["q"], + k=a["k"], + v=a["v"], + g=a["g"], + beta=a["beta"], + scale=None, + initial_state=pool, + initial_state_indices=a["initial_state_indices"], + inplace_indexed_state_update=True, + output_final_state=False, + cu_seqlens=a["cu_seqlens"], + head_first=False, + use_qk_l2norm_in_kernel=False, + output=None, + ) + + warm = _build(dict(geometry=QWEN38_GEOMETRY, seqlens=(64,), slots=1, indices=(0,)), seed=1) + run_triton(warm, warm["initial_state"]) + args = _build( + dict(geometry=QWEN38_GEOMETRY, seqlens=(1024, 37, 1), slots=461, indices=(460, 0, 17)), + seed=3, + ) + pool_t = args["initial_state"].clone() + o_t, _ = run_triton(args, pool_t) + o, _ = op( + args["q"], + args["k"], + args["v"], + args["g"], + args["beta"], + None, + args["initial_state"], + args["initial_state_indices"], + True, + False, + args["cu_seqlens"], + False, + False, + None, + ) + torch.testing.assert_close(o.float(), o_t.float(), rtol=RTOL, atol=ATOL_OUTPUT) + torch.testing.assert_close(args["initial_state"], pool_t, rtol=RTOL, atol=ATOL_STATE) diff --git a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_gated_delta_rule_decode_pretranspose.py b/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_gated_delta_rule_decode_pretranspose.py new file mode 100644 index 000000000000..77bc3fd1732e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_gated_delta_rule_decode_pretranspose.py @@ -0,0 +1,167 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gated_delta_rule_decode_pretranspose catalog entry (FlashInfer fp32-pool decode).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.gated_delta_rule_decode_pretranspose import ( # noqa: E501 + QWEN38_GEOMETRY, + _scattered, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.gated_delta_rule_decode_pretranspose import ( # noqa: E501 + gated_delta_rule_decode_pretranspose as op, +) + +__extra_import_path__ = [".."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "gated_delta_rule_decode_pretranspose requires a CUDA device" + +DEV = "cuda" +_PAD = 64 + + +def _compare(out, ref) -> None: + """`(o, pool)` against the fp64 reference rounded once, at torch's default + dtype-aware band. Measured on the microbench inputs: o max abs 1.2e-4 at + |o| <= 6.2e-2 (identical to the Triton op's), state max abs 2.6e-7. The + pool is compared whole, so slots the call did not name are gated too.""" + (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.shape == ref_pool.shape and pool.dtype == ref_pool.dtype + torch.testing.assert_close(pool, ref_pool) + + +def _build(spec, seed): + """`__call__`'s arguments (by name) plus the buffers some layouts view.""" + g = torch.Generator(device=DEV).manual_seed(seed) + H, HV, K, V = spec["geometry"] + m, slots = spec["m"], spec["slots"] + mixed = torch.randn(m, 2 * H * K + HV * V, generator=g, device=DEV, dtype=torch.bfloat16) + qv, kv, vv = torch.split(mixed, [H * K, H * K, HV * V], 1) + q, k, v = ( + torch.reshape(qv, [m, 1, H, K]), + torch.reshape(kv, [m, 1, H, K]), + torch.reshape(vv, [m, 1, HV, V]), + ) + if spec.get("qkv", "views") == "contiguous": + q, k, v = q.contiguous(), k.contiguous(), v.contiguous() + a = torch.randn(m, HV, generator=g, device=DEV, dtype=torch.bfloat16).view(m, 1, HV) + b = torch.randn(m, HV, generator=g, device=DEV, dtype=torch.bfloat16).view(m, 1, HV) + A_log = torch.empty(HV, device=DEV).uniform_(-5.5, 1.0, generator=g) + dt_bias = torch.randn(HV, generator=g, device=DEV) * 4.0 + layout = spec.get("pool", "contiguous") + extras = {} + if layout == "padded": + backing = torch.randn(slots, HV * V * K + _PAD, generator=g, device=DEV) * 0.05 + pool = backing[:, : HV * V * K].view(slots, HV, V, K) + extras["backing"] = backing + elif layout == "layer_view": + cache = torch.randn(3, slots, HV, V, K, generator=g, device=DEV) * 0.05 + pool = cache[1] + extras["cache"] = cache + else: + pool = torch.randn(slots, HV, V, K, generator=g, device=DEV) * 0.05 + indices = spec.get("indices") or _scattered(m, slots) + indices = torch.tensor(indices, dtype=torch.int32, device=DEV) + output = None + if spec.get("output", "buffer") == "buffer": + output = torch.full((m, 1, HV, V), -7.0, device=DEV, dtype=torch.bfloat16) + args = dict( + q=q, + k=k, + v=v, + state=None, + A_log=A_log, + a=a, + dt_bias=dt_bias, + b=b, + scale=None, + output=output, + use_qk_l2norm=True, + initial_state=pool, + initial_state_indices=indices, + output_state_indices=None, + ) + 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) + args, extras = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) + pool = args["initial_state"] + pool_before = pool.clone() + cache_before = extras["cache"].clone() if "cache" in extras else None + slot_numel = pool[0].numel() + pad_before = extras["backing"][:, slot_numel:].clone() if "backing" in extras else None + ref = op.reference(**args) + + o, returned_pool = op(**args) + + assert returned_pool.data_ptr() == pool.data_ptr(), "the pool is advanced in place" + _compare((o, pool), ref) + if args["output"] is not None: + assert o.data_ptr() == args["output"].data_ptr(), "output= must be written in place" + named = {s for s in args["initial_state_indices"].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" + for slot in named: + assert not torch.equal(pool[slot], pool_before[slot]), slot + if cache_before is not None: + for other in (0, 2): + assert torch.equal(extras["cache"][other], cache_before[other]), other + if pad_before is not None: + assert torch.equal(extras["backing"][:, slot_numel:], pad_before), "slot padding written" + + +def test_repeated_calls_are_bitwise() -> None: + """Two identical calls from the same start give identical output and state.""" + results = [] + for _ in range(2): + args, _ = _build(dict(geometry=QWEN38_GEOMETRY, m=8, slots=40), seed=5) + o, pool = op(**args) + results.append((o.clone(), pool.clone())) + assert torch.equal(results[0][0], results[1][0]) and torch.equal(results[0][1], results[1][1]) + + +def test_matches_in_tree_triton_op() -> None: + """Same inputs through the in-tree Triton decode op this entry stands in + for: output and state agree at the default band (each is within it of fp64).""" + from tensorrt_llm._torch.modules.fla.fused_sigmoid_gating_recurrent import ( + fused_sigmoid_gating_delta_rule_update as triton_op, + ) + + H, HV, K, V = QWEN38_GEOMETRY + m = 32 + args, _ = _build(dict(geometry=QWEN38_GEOMETRY, m=m, slots=461), seed=11) + pool_t = args["initial_state"].clone() + cu = torch.arange(m + 1, device=DEV, dtype=torch.int64) + o_t = triton_op( + A_log=args["A_log"], + a=args["a"].view(m, HV), + dt_bias=args["dt_bias"], + softplus_beta=1.0, + softplus_threshold=20.0, + q=torch.reshape(args["q"], [1, m, H, K]), + k=torch.reshape(args["k"], [1, m, H, K]), + v=torch.reshape(args["v"], [1, m, HV, V]), + b=args["b"].view(m, HV), + initial_state_source=pool_t, + initial_state_indices=args["initial_state_indices"], + scale=None, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cu, + output=None, + ) + o, pool = op(**args) + torch.testing.assert_close(o.view(1, m, HV, V), o_t) + torch.testing.assert_close(pool, pool_t) From acafef16a17efaff493f01970ca547d5b905649c Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Sun, 18 Oct 2026 23:00:00 -0700 Subject: [PATCH 05/28] [TRTLLM-16460][perf] Qwen3.8-Flash-Next target: attention ops The target's QSA attention fast paths over a bf16 pool, with their unit tests: - qsa_sparse_gqa_rowtile + qsa_rowtile_attention_cute: context batches, several query rows a program with each K/V tile read once; the row-tile pass on SM103 tensor cores (CuTe DSL, tcgen05). - qsa_paged_sparse_gqa_decode: decode rows, a CuTe DSL split-K pass that loads its tokens in one round trip, then a merge spread over the SMs, both under PDL. - qwen4_exp_attn_decode_prologue: a sparse decode layer's work between its input projection and its index scores in one PDL launch; qsa_topk_expand_decode: the block top-k and its expansion to token indices in one launch. - qwen4_exp_qk_norm_rope: Gemma QK norm + partial NeoX rope for both planes in one launch (a persistent CuTe ring at prefill sizes). - Updates to qsa_paged_index_scores / _index_store / _kv_store / _sparse_gqa / qsa_topk_select for the calls above (PDL, row strides, cells). Folds these commits of the review branch (their changes to the files of this commit): - qwen3.8 target: carry the decode-fusion perf work onto the feature target (62151b727b) - qwen3.8 target: canonical decode selection in one launch (ac8ad93f0f) - qwen3.8 target: canonical context-row selection in one launch (bd29d1c90f) - qwen3.8 target: row-tile bitmap pass without global atomics (3861769b4f) - qwen3.8 target: three CTAs an SM in the block top-k + expansion (f181186145) - qwen3.8 target: QK norm + rope on a persistent CuTe ring at prefill sizes (d865464dca) Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../ops/attention/qsa_paged_index_scores.py | 33 +- .../ops/attention/qsa_paged_index_store.py | 438 ++++-- .../ops/attention/qsa_paged_kv_store.py | 11 + .../ops/attention/qsa_paged_sparse_gqa.py | 72 +- .../attention/qsa_paged_sparse_gqa_decode.py | 828 ++++++++++++ .../attention/qsa_rowtile_attention_cute.py | 998 ++++++++++++++ .../ops/attention/qsa_sparse_gqa_rowtile.py | 1043 ++++++++++++++ .../ops/attention/qsa_topk_expand_decode.py | 590 ++++++++ .../ops/attention/qsa_topk_select.py | 50 +- .../qwen4_exp_attn_decode_prologue.py | 576 ++++++++ .../ops/attention/qwen4_exp_qk_norm_rope.py | 1201 +++++++++++++++++ ...test_modeling_v2_qsa_paged_index_scores.py | 74 +- .../test_modeling_v2_qsa_paged_index_store.py | 190 ++- .../test_modeling_v2_qsa_paged_kv_store.py | 8 +- .../test_modeling_v2_qsa_paged_sparse_gqa.py | 50 +- ...modeling_v2_qsa_paged_sparse_gqa_decode.py | 163 +++ ...test_modeling_v2_qsa_sparse_gqa_rowtile.py | 468 +++++++ ...test_modeling_v2_qsa_topk_expand_decode.py | 164 +++ ...eling_v2_qwen4_exp_attn_decode_prologue.py | 131 ++ ...test_modeling_v2_qwen4_exp_qk_norm_rope.py | 363 +++++ 20 files changed, 7206 insertions(+), 245 deletions(-) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_rowtile_attention_cute.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_expand_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_attn_decode_prologue.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_qk_norm_rope.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_decode.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_expand_decode.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qwen4_exp_attn_decode_prologue.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qwen4_exp_qk_norm_rope.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_scores.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_scores.py index 5c626cf40ae7..6675d234850f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_scores.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_scores.py @@ -98,9 +98,9 @@ class QsaPagedIndexScores(OpWrapper): `context_rows` is a layout declaration ("consecutive rows mostly share a request"), not a phase one: with more than 256 rows and an index-head count - that divides 16 into more than one row (1, 2, 4, 8) it selects a - query-tiled kernel sharing one key gather across rows; otherwise the - row-wise kernel runs. Each cell names the kernel it was measured on. The + that divides 128 into more than one row (1, 2, 4, ..., 64) it selects a + query-tiled kernel sharing one key gather across `128 // heads` rows; + otherwise the row-wise kernel runs. Each cell names the kernel it was measured on. The two agree to one fp32 ulp, not bitwise. Fusion boundary: the table walk, the gather, the dots, relu, head sum, @@ -389,27 +389,27 @@ class QsaPagedIndexScores(OpWrapper): ), 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", + why="300 rows x 1 head: 128 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", + why="300 rows x 2 heads: 64 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", + why="300 rows x 8 heads: 16 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), + why="300 rows x 16 heads: 8 rows per query tile", + spec=_small(300, 5, heads=16, context_rows=True, branch="query_tile"), ), Cell( - why="300 rows x 3 heads: does not divide 16, row-wise", + why="300 rows x 3 heads: does not divide 128, row-wise", spec=_small(300, 5, heads=3, context_rows=True), ), Cell( - why="300 rows x 5 heads: does not divide 16, row-wise", + why="300 rows x 5 heads: does not divide 128, row-wise", spec=_small(300, 5, heads=5, context_rows=True), ), Cell( @@ -474,6 +474,13 @@ class QsaPagedIndexScores(OpWrapper): again after `q`, the cache and the table were overwritten in place; both matched the reference. + `prefetch` (with `use_pdl` and `only_visible_blocks`, at generation-sized + batches) reads the pages and compressed keys before the PDL wait and q after + it. Sound only behind the decode step's own prologue + (`qwen4_exp_attn_decode_prologue`), which writes no compressed key but the last + visible block's, at a row whose position ends a group; that key is read again + after the wait. + 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 @@ -491,6 +498,8 @@ def __call__( compress_ratio: int, only_visible_blocks: bool, context_rows: bool, + use_pdl: bool = False, + prefetch: bool = False, ) -> torch.Tensor: return _triton_qsa_paged_index_scores( q=q, @@ -502,6 +511,8 @@ def __call__( compress_ratio=compress_ratio, only_visible_blocks=only_visible_blocks, context_rows=context_rows, + use_pdl=use_pdl, + prefetch=prefetch, ) def reference( @@ -515,6 +526,8 @@ def reference( compress_ratio: int, only_visible_blocks: bool, context_rows: bool, + use_pdl: bool = False, + prefetch: bool = False, ) -> torch.Tensor: """fp32 scores, `-inf` past each row's bound -- or NaN there, meaning "unspecified", when `only_visible_blocks` is True. `context_rows` is diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_store.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_store.py index bef6cf362071..918cb5f91716 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_store.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_store.py @@ -7,9 +7,114 @@ from typing import Tuple import torch +import triton +import triton.language as tl from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +_INTEGER_DTYPES = (torch.int32, torch.int64) + + +# The row count and the pool/table extents change from step to step; letting +# Triton specialize on them (== 1, divisible by 16) only multiplies compiles. +@triton.jit(do_not_specialize=["num_rows", "num_requests", "table_width", "num_pages"]) +def _qsa_paged_index_store_kernel( + token_k, + coords, + index_cache, + position_cache, + request_indices, + logical_positions, + block_table, + num_rows, + num_requests, + table_width, + num_pages, + stride_k_row, + stride_k_head, + stride_k_dim, + stride_c_row, + stride_c_axis, + stride_i_page, + stride_i_slot, + stride_i_head, + stride_i_dim, + stride_p_page, + stride_p_slot, + stride_p_axis, + stride_request, + stride_position, + stride_t_request, + stride_t_column, + HKV: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, + AXES: tl.constexpr, + BLOCK_A: tl.constexpr, + TOKENS_PER_BLOCK: tl.constexpr, + BLOCK_ROWS: tl.constexpr, +): + # One program stores BLOCK_ROWS rows: the page lookup, then every head of + # the key and the coordinate row, all through the operands' own strides. + rows = (tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)).to(tl.int64) + live = rows < num_rows + request = tl.load(request_indices + rows * stride_request, mask=live, other=-1).to(tl.int64) + position = tl.load(logical_positions + rows * stride_position, mask=live, other=-1).to(tl.int64) + column = position // TOKENS_PER_BLOCK + slot = position % TOKENS_PER_BLOCK + live &= (request >= 0) & (request < num_requests) & (position >= 0) & (column < table_width) + page = tl.load( + block_table + request * stride_t_request + column * stride_t_column, mask=live, other=-1 + ).to(tl.int64) + live &= (page >= 0) & (page < num_pages) + + dims = tl.arange(0, BLOCK_D) + key_mask = live[:, None] & (dims < HEAD_DIM)[None, :] + for head in tl.static_range(HKV): + key = tl.load( + token_k + + rows[:, None] * stride_k_row + + head * stride_k_head + + dims[None, :] * stride_k_dim, + mask=key_mask, + other=0.0, + ) + tl.store( + index_cache + + page[:, None] * stride_i_page + + slot[:, None] * stride_i_slot + + head * stride_i_head + + dims[None, :] * stride_i_dim, + key.to(index_cache.dtype.element_ty), + mask=key_mask, + ) + + axes = tl.arange(0, BLOCK_A) + axis_mask = live[:, None] & (axes < AXES)[None, :] + coordinate = tl.load( + coords + rows[:, None] * stride_c_row + axes[None, :] * stride_c_axis, + mask=axis_mask, + other=0, + ) + tl.store( + position_cache + + page[:, None] * stride_p_page + + slot[:, None] * stride_p_slot + + axes[None, :] * stride_p_axis, + coordinate.to(position_cache.dtype.element_ty), + mask=axis_mask, + ) + + +def _launch_rows(rows: int) -> Tuple[int, int]: + """Rows per program and warps: one row per program while the grid is small, + a tile of rows once there are enough rows to fill the machine. Measured on + GB300 against 1/4/8/16/32 rows per program (microbench `qsa-index-store-kernel`).""" + if rows <= 256: + return 1, 1 + return 16, 4 + + #: This checkpoint's index geometry and the engine's QSA metadata. `token_k` is #: the k head of the packed `[T, 4 + 1, 128]` index projection; the coordinates #: are the `[T, 3]` transpose of the engine's `[3, max_num_tokens]` position @@ -30,10 +135,19 @@ meta_layout="contiguous", cache_layout="plain", table_layout="dense", + hole=None, ) _DECODE = ((0, 1), (1, 1), (31, 1), (32, 1), (33, 1), (63, 1), (64, 1), (95, 1)) +#: The campaign's envelope: ISL 8192 + OSL 1024 puts decode rows up to position +#: 9215 (page column 287), 128 concurrent requests, and `max_num_tokens` 16384 +#: rows in one prefill or mixed step. +_WIDE = dict(slots=160, width=300) +_DECODE_128 = tuple(((r * 72 + 9) % 9216, 1) for r in range(128)) +_PREFILL_16K = ((0, 8192), (0, 8192)) +_MIXED_16K = ((0, 8192), (1024, 8064)) + tuple(((r * 71 + 5) % 9216, 1) for r in range(128)) + #: 32 requests of one MTP7 verify block each (8 rows), at starts that put some blocks across a #: page edge and some inside one. _MTP7_32 = tuple((29 + 61 * i, 8) for i in range(32)) @@ -44,11 +158,6 @@ def _spec(segments, pages, **overrides) -> dict: return {**_CKPT, "segments": segments, "pages": pages, **overrides} -def _bits(tensor: torch.Tensor) -> torch.Tensor: - as_int = {1: torch.uint8, 2: torch.int16, 4: torch.int32}[tensor.element_size()] - return tensor.view(as_int) - - class QsaPagedIndexStore(OpWrapper): """Write each row's raw index key and position coordinates into its paged slot. @@ -58,16 +167,21 @@ class QsaPagedIndexStore(OpWrapper): index_cache[page, slot] = token_k[t] # cast to the cache dtype position_cache[page, slot, :A] = position_coordinates[t] # A = coordinates' width - The key is stored raw -- un-normed, un-rotated: `qsa_prefill_compress` - pools raw keys and applies the Gemma norm and the rotation itself. Returns - `(index_cache, position_cache)`, both written in place. - - 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. + A row that resolves to no slot -- negative request, request past the + table, negative position, page column past the table, or a table entry + that is negative (`-1`, unallocated) or past the pool -- is skipped + without error. The key is stored raw -- un-normed, un-rotated: + `qsa_prefill_compress` pools raw keys and applies the Gemma norm and the + rotation itself. Returns `(index_cache, position_cache)`, both written in + place. + + Sourcing tier 5: one Triton launch written for this target + (`_qsa_paged_index_store_kernel`). Nothing in the engine performs this + store for a packed, ragged batch: the registered index-cache ops write fp8 + planes or norm and rotate on the way in, `triton_qsa_decode_pre_indexer` + is one row per request, and `QSAIndexer.update_cache_and_compress` bundles + the compression. Native torch needs a block-table gather, a floor-divide, + a remainder, two indexed writes and their index casts (7 launches). Fusion boundary: page lookup, slot arithmetic, both scatters. The index projection, the query's norm and rotation, group compression, scoring, @@ -101,6 +215,14 @@ class QsaPagedIndexStore(OpWrapper): 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="decode: token_k a head of the fused projection row (110 x 128)", + spec=_spec(_DECODE, pages=32, key_layout="fused_row"), + ), + Cell( + why="prefill 1024 rows: token_k off the fused projection row", + spec=_spec(((0, 256),) * 4, pages=32, key_layout="fused_row"), + ), Cell( why="contiguous token_k and coordinates", spec=_spec(_DECODE, pages=32, key_layout="contiguous", coords_layout="contiguous"), @@ -137,6 +259,10 @@ class QsaPagedIndexStore(OpWrapper): 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="block table transposed: column stride != 1", + spec=_spec(_DECODE, pages=32, table_layout="transposed"), + ), Cell( why="fp16 keys: one conversion into the bf16 plane", spec=_spec(((0, 4),), pages=8, key_dtype=torch.float16), @@ -153,6 +279,56 @@ class QsaPagedIndexStore(OpWrapper): why="int64 requests and int64 positions", spec=_spec(_DECODE, pages=32, request_dtype=torch.int64), ), + # ---- the campaign envelope: decode M=128, 16384-row steps ---------- + Cell( + why="decode M=128 at positions up to 9153 (page column 286)", + spec=_spec(_DECODE_128, pages=18432, **_WIDE), + ), + Cell( + why="decode M=128, token_k off the fused projection row", + spec=_spec(_DECODE_128, pages=18432, key_layout="fused_row", **_WIDE), + ), + Cell( + why="prefill 16384 rows: two 8192-token requests (max_num_tokens)", + spec=_spec(_PREFILL_16K, pages=600, **_WIDE), + ), + Cell( + why="prefill 16384 rows, token_k off the fused projection row", + spec=_spec(_PREFILL_16K, pages=600, key_layout="fused_row", **_WIDE), + ), + Cell( + why="mixed 16384 rows: two chunks, one mid-sequence, + 128 decode rows", + spec=_spec(_MIXED_16K, pages=18688, **_WIDE), + ), + Cell( + why="head_dim 64, 2 index-KV heads, 16-token pages", + spec=_spec(_DECODE, pages=48, hkv=2, head_dim=64, tokens_per_block=16), + ), + # ---- rows that resolve to no slot are skipped ------------------------ + Cell( + why="skip: table entry -1 (unallocated)", + spec=_spec(_DECODE, pages=32, hole="unallocated_page"), + ), + Cell( + why="skip: table entry past the pool", + spec=_spec(_DECODE, pages=32, hole="page_past_pool"), + ), + Cell( + why="skip: request index past the table", + spec=_spec(_DECODE, pages=32, hole="request_past_table"), + ), + Cell( + why="skip: negative request index", + spec=_spec(_DECODE, pages=32, hole="negative_request"), + ), + Cell( + why="skip: negative logical position", + spec=_spec(_DECODE, pages=32, hole="negative_position"), + ), + Cell( + why="skip: page column past the table", + spec=_spec(_DECODE, pages=32, hole="column_past_table"), + ), Cell(why="zero rows: both planes unchanged", spec=_spec((), pages=8)), ) @@ -163,39 +339,49 @@ class QsaPagedIndexStore(OpWrapper): `get_index_k_buffer(i)` (`[pages, tpb, 1, 128]` bf16) and the shared `get_qsa_position_buffer()` (`[pages, tpb, 3]` int32, rewritten idempotently by every sparse layer). Every operand is read through its - own strides, unlike the Triton siblings. - - Values the caller owns -- data, not metadata, so not guarded; each - measured, and unlike `qsa_paged_kv_store` nothing is skipped: - - a `-1` (unallocated) table entry wraps to the pool's last page, and a - negative request selects the table's last row; - - a negative position wraps twice (last table column, last slot); - - a page, request, column or slot genuinely out of range is a CUDA - device-side assert that poisons the process -- so none is tested; - - duplicate destinations both store, unordered; - - coordinates wider than the plane's dtype wrap (int64 `2^31 + 5` - stores as -2147483643); floats truncate toward zero. - Under an NVFP4 KV pool the QSA cache manager packs each page's index-K - roles and then its position page into one per-slot buffer, so both planes - are strided views of one allocation: the trunk's index plane is - `[pages, 32, 1, 128]` at page stride 49344 (12 x 4096 + 192) with the - position plane `[pages, 32, 3]` at (24672, 3, 1); the cell checks that the - other index-K roles of each page survive the call. - - 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 (captured at 4 decode rows and at 256, 32 MTP7 - verify blocks); a replay follows positions rewritten in place. Repeating a - call is idempotent. - - Hard limits: torch advanced indexing in int64 -- no int32 boundary inside - the call; the coordinate wrap above is the plane's dtype. - - Not certified: an fp8 index plane, `Hkv` above 1, multiple devices, - concurrent streams onto one pool. + own strides, all of them runtime arguments: `token_k` with any row stride + (the packed index projection's 640, the fused input projection's 14080, + contiguous), stride-0 or transposed coordinates, strided metadata, a + column-sliced or transposed table, a role view of the index plane. Under an + NVFP4 KV pool the QSA cache manager packs each page's index-K roles and + then its position page into one per-slot buffer, so both planes are strided + views of one allocation: the trunk's index plane is `[pages, 32, 1, 128]` + at page stride 49344 (12 x 4096 + 192) with the position plane + `[pages, 32, 3]` at (24672, 3, 1); the cell checks that the other index-K + roles of each page survive the call. + + Preconditions: + - destination slots are unique across the batch: two rows naming one + (page, slot) both store, unordered. The engine's mapping cannot + collide. + Loud rejections (`ValueError`, before launching): metadata that is not + one-dimensional, of unequal length, or not int32/int64; `token_k` or the + coordinates with a row count other than the metadata's; an index plane + whose page size, head count or head width disagrees with + `tokens_per_block` and `token_k`; a position plane whose page size + disagrees or that is narrower than the coordinates; a table that is not a + 2-D integer tensor. `rows == 0` returns before launching. + + Values the caller owns -- data, not metadata: coordinates wider than the + plane's dtype wrap (int64 `2^31 + 5` stores as -2147483643); floats + truncate toward zero. + + One launch per call; HEAD_DIM, the head count, the page size and the + coordinate width are constexprs, everything else (strides, pool and table + extents, rows) is a runtime argument and the row count and extents are + not specialized on, so a new batch or pool size does not recompile + (strides still specialize on 1 and 16-divisibility). Up to 256 rows run + one row per program (1 warp), above that 16 rows per program (4 warps). + CUDA-graph capturable (captured at 4 decode rows and at 256, 32 MTP7 verify + blocks); a replay follows positions rewritten in place. Repeating a call is + idempotent. + + Hard limits: every index is widened to int64 inside the kernel before + any arithmetic; an int64 position at `2^32 + 5` lands past the table and + is skipped rather than truncated. Rows are a grid axis: 2^31 programs. + + Not certified: an fp8 index plane, multiple devices, concurrent streams + onto one pool, duplicate destinations. """ def __call__( @@ -209,14 +395,73 @@ def __call__( block_table: torch.Tensor, tokens_per_block: int, ) -> Tuple[torch.Tensor, torch.Tensor]: - positions = logical_positions.to(torch.long) - pages = block_table[request_indices.to(torch.long), positions // tokens_per_block].to( - torch.long - ) - within = positions % tokens_per_block - index_cache[pages, within] = token_k.to(index_cache.dtype) - position_cache[pages, within, : position_coordinates.shape[-1]] = position_coordinates.to( - position_cache.dtype + rows = request_indices.shape[0] + if request_indices.dim() != 1 or logical_positions.shape != request_indices.shape: + raise ValueError( + f"request_indices {tuple(request_indices.shape)} and logical_positions " + f"{tuple(logical_positions.shape)} must be one-dimensional and equally long" + ) + if request_indices.dtype not in _INTEGER_DTYPES or ( + logical_positions.dtype not in _INTEGER_DTYPES + or block_table.dtype not in _INTEGER_DTYPES + ): + raise ValueError("QSA index store metadata and block table must be int32 or int64") + if token_k.dim() != 3 or token_k.shape[0] != rows: + raise ValueError(f"token_k {tuple(token_k.shape)} must be [{rows}, heads, dim]") + if position_coordinates.dim() != 2 or position_coordinates.shape[0] != rows: + raise ValueError( + f"position_coordinates {tuple(position_coordinates.shape)} must be [{rows}, axes]" + ) + if index_cache.dim() != 4 or index_cache.shape[1:] != ( + tokens_per_block, + *token_k.shape[1:], + ): + raise ValueError( + f"index_cache {tuple(index_cache.shape)} must be [pages, {tokens_per_block}, " + f"{token_k.shape[1]}, {token_k.shape[2]}]" + ) + if ( + position_cache.dim() != 3 + or position_cache.shape[1] != tokens_per_block + or position_cache.shape[2] < position_coordinates.shape[1] + ): + raise ValueError( + f"position_cache {tuple(position_cache.shape)} must be [pages, " + f"{tokens_per_block}, >= {position_coordinates.shape[1]}]" + ) + if block_table.dim() != 2: + raise ValueError(f"block_table {tuple(block_table.shape)} must be two-dimensional") + if rows == 0: + return index_cache, position_cache + block_rows, warps = _launch_rows(rows) + axes = position_coordinates.shape[1] + _qsa_paged_index_store_kernel[(triton.cdiv(rows, block_rows),)]( + token_k, + position_coordinates, + index_cache, + position_cache, + request_indices, + logical_positions, + block_table, + rows, + block_table.shape[0], + block_table.shape[1], + index_cache.shape[0], + *token_k.stride(), + *position_coordinates.stride(), + *index_cache.stride(), + *position_cache.stride(), + request_indices.stride(0), + logical_positions.stride(0), + *block_table.stride(), + HKV=token_k.shape[1], + HEAD_DIM=token_k.shape[2], + BLOCK_D=triton.next_power_of_2(token_k.shape[2]), + AXES=axes, + BLOCK_A=triton.next_power_of_2(axes), + TOKENS_PER_BLOCK=tokens_per_block, + BLOCK_ROWS=block_rows, + num_warps=warps, ) return index_cache, position_cache @@ -231,70 +476,39 @@ def reference( block_table: torch.Tensor, tokens_per_block: int, ) -> Tuple[torch.Tensor, torch.Tensor]: - """Both planes after the store, as fresh tensors, built row by row. + """Both planes after the store, as fresh tensors; the inputs are not touched. - Pages and slots are computed in Python and each write is a basic - (integer) index -- none of the advanced indexing `__call__` is made of. + Every row's destination, and whether it has one, is decided in Python + integers from the table read back to the host -- none of the kernel's + arithmetic -- and the surviving rows are written with one indexed copy + per plane. """ expected_index, expected_position = index_cache.clone(), position_cache.clone() table = block_table.tolist() - axes = position_coordinates.shape[-1] + pool = index_cache.shape[0] + kept, pages, slots = [], [], [] rows = zip(request_indices.tolist(), logical_positions.tolist()) for row, (request, position) in enumerate(rows): - page = table[request][position // tokens_per_block] - slot = position % tokens_per_block - expected_index[page, slot] = token_k[row].to(index_cache.dtype) - expected_position[page, slot, :axes] = position_coordinates[row].to( - position_cache.dtype - ) - return expected_index, expected_position - - def is_valid( - self, - token_k: torch.Tensor, - position_coordinates: torch.Tensor, - index_cache: torch.Tensor, - position_cache: torch.Tensor, - request_indices: torch.Tensor, - logical_positions: torch.Tensor, - block_table: torch.Tensor, - tokens_per_block: int, - ) -> None: - # Each is a shape check over a violation measured to be accepted and - # silently wrong: torch broadcasts instead of raising. - rows = request_indices.shape[0] - assert request_indices.dim() == 1 and logical_positions.dim() == 1, ( - f"request_indices {tuple(request_indices.shape)} and logical_positions " - f"{tuple(logical_positions.shape)} must both be one-dimensional; a " - f"trailing axis broadcasts into a rows x rows index instead of raising" - ) - assert logical_positions.shape[0] == rows, ( - f"logical_positions has {logical_positions.shape[0]} entries against " - f"{rows} request_indices; a length-1 operand broadcasts silently and a " - f"length-1 logical_positions collapses every row onto one slot" - ) - assert token_k.shape[0] == rows, ( - f"token_k has {token_k.shape[0]} rows against {rows} metadata rows; a " - f"single row broadcasts into every destination slot instead of raising" - ) - assert position_coordinates.shape[0] == rows, ( - f"position_coordinates has {position_coordinates.shape[0]} rows against " - f"{rows} metadata rows; a single row broadcasts into every slot" - ) - assert index_cache.shape[1] == tokens_per_block, ( - f"tokens_per_block={tokens_per_block} does not match index_cache's " - f"token axis {index_cache.shape[1]}; a wrong page size writes the " - f"wrong page and the wrong slot" + if not 0 <= request < len(table) or position < 0: + continue + column, slot = divmod(position, tokens_per_block) + if column >= len(table[request]) or not 0 <= table[request][column] < pool: + continue + kept.append(row) + pages.append(table[request][column]) + slots.append(slot) + if not kept: + return expected_index, expected_position + device = index_cache.device + kept_t, page_t, slot_t = ( + torch.tensor(v, dtype=torch.long, device=device) for v in (kept, pages, slots) ) - assert position_cache.shape[1] == tokens_per_block, ( - f"tokens_per_block={tokens_per_block} does not match position_cache's " - f"token axis {position_cache.shape[1]}" - ) - assert index_cache.shape[2] == token_k.shape[1], ( - f"index_cache carries {index_cache.shape[2]} index-KV heads against " - f"token_k's {token_k.shape[1]}; a one-head token_k broadcasts into " - f"every head of a wider cache" + axes = position_coordinates.shape[-1] + expected_index[page_t, slot_t] = token_k[kept_t].to(index_cache.dtype) + expected_position[page_t, slot_t, :axes] = position_coordinates[kept_t].to( + position_cache.dtype ) + return expected_index, expected_position qsa_paged_index_store = QsaPagedIndexStore() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store.py index 8caa29166a2b..e532a87f785f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store.py @@ -34,6 +34,9 @@ k_layout="qkv_slice", table_layout="dense", hole=None, + # Heads of further columns after v in each row: 27 is the target's fused input + # projection row [q | k | v | gate | index | pad] (55 heads of 256). + row_pad_heads=0, values="randn", ) @@ -138,6 +141,14 @@ class QsaPagedKvStore(OpWrapper): 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="decode: k/v as head slices of the fused projection row (row stride 14080)", + spec=_spec(_DECODE, pages=64, row_pad_heads=27), + ), + Cell( + why="prefill 1024 rows off the fused projection row (row stride 14080)", + spec=_spec(((0, 256),) * 4, pages=160, row_pad_heads=27), + ), Cell( why="head-major rows: k/v last-dim stride != 1", spec=_spec(_DECODE, pages=64, k_layout="head_major"), diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa.py index d6ebc974a7d9..be7d55e4991b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa.py @@ -12,6 +12,9 @@ # Approved exception: QSA's sparse path never reaches the thop attention op's # sparse argument group, so there is no torch.ops surface for this computation. +from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import ( + can_fuse_qsa_splitk_output_gate as _can_fuse_qsa_splitk_output_gate, +) from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import ( triton_qsa_paged_sparse_gqa as _triton_qsa_paged_sparse_gqa, ) @@ -27,7 +30,8 @@ q_dtype=torch.bfloat16, # K and V are the two role planes of one [pages, 2, kv, tpb, D] pool. cache_layout="pool_planes", - # q is the first 24 heads of the packed [T, 24 + 2 + 2, 256] QKV buffer. + # q is the first 24 heads of the packed [T, 24 + 2 + 2, 256] QKV buffer; + # "fused_row" is the target's fused input projection row, 55 heads of 256. q_layout="qkv_slice", request_dtype=torch.int32, # The engine's table: 260 pages per slot, -1 past each sequence. @@ -91,9 +95,20 @@ class QsaPagedSparseGqa(OpWrapper): silently truncated (measured: a planted index moving the right answer by 0.77 moved the op's by nothing). The split-K kernel ignores the pair. - Fusion boundary: the block-table gather, scale, online softmax, P.V and - the split-K merge are inside. Deciding what to attend to, appending K/V to - the pool, and q's projection, norm, RoPE and output gate are not. + **The output gate rides the merge when `output_gate` is passed.** Each + merged head row is rounded to `q.dtype`, multiplied in fp32 by + `sigmoid(output_gate)` and rounded again: the arithmetic of the ungated + call followed by `fused_sigmoid_mul_inplace`. `output_gate` is `[rows, + Hq * D]` or `[rows, Hq, D]` in `q.dtype`, unit stride in its last dim, any + row stride (the gate columns of the fused projection row). Only a split-K + launch with more than one split has a merge to carry it; `fuses_output_gate` + says whether a call does, and the op raises `ValueError` if handed a gate + it cannot carry. + + Fusion boundary: the block-table gather, scale, online softmax, P.V, the + split-K merge and, when passed, the output gate are inside. Deciding what + to attend to, appending K/V to the pool, and q's projection, norm and RoPE + are not. """ ARCHS = frozenset({Arch.SM_103}) @@ -108,6 +123,22 @@ class QsaPagedSparseGqa(OpWrapper): why="decode: 4 requests x 1 row, split-K + merge, PDL", spec=_decode(((4000, 1), (4037, 1), (4074, 1), (4111, 1))), ), + Cell( + why="decode: 1 row, gate off the fused row, in the merge, PDL", + spec=_decode(((4000, 1),), q_layout="fused_row", gate="fused_row"), + ), + Cell( + why="decode: 4 requests x 1 row, gate fused into the merge", + spec=_decode(tuple((4000 + 37 * i, 1) for i in range(4)), gate="fused_row"), + ), + Cell( + why="decode: 128 rows, 8-way split-K, gate fused into the merge", + spec=_decode(((600, 1),) * 128, q_layout="fused_row", slots=160, gate="fused_row"), + ), + Cell( + why="decode: 4 rows, a compact [rows, Hq, D] gate in the merge", + spec=_decode(((900, 1),) * 4, gate="heads"), + ), Cell( why="decode MTP: 4 requests x 2 rows, split-K + merge", spec=_decode(tuple((4000 + 37 * i, 2) for i in range(4))), @@ -220,6 +251,18 @@ class QsaPagedSparseGqa(OpWrapper): Cell( why="int64 request indices", spec=_decode(((1000, 2),) * 4, request_dtype=torch.int64) ), + Cell( + why="q as the head slice of the fused projection row (55 x 256), decode", + spec=_decode(((900, 1),) * 4, q_layout="fused_row"), + ), + Cell( + why="q off the fused projection row (55 x 256), 128 decode rows", + spec=_decode(((600, 1),) * 128, q_layout="fused_row", slots=160), + ), + Cell( + why="q off the fused projection row (55 x 256), 8191-row prefill", + spec=_decode(((0, 8191),), launches=_PREFILL, prefill_pair=True, q_layout="fused_row"), + ), Cell( why="standalone contiguous K and V buffers, contiguous q", spec=_decode(((900, 1),) * 4, cache_layout="contiguous", q_layout="contiguous"), @@ -286,6 +329,7 @@ def __call__( softmax_scale: float, query_positions: Optional[torch.Tensor] = None, compress_ratio: Optional[int] = None, + output_gate: Optional[torch.Tensor] = None, ) -> torch.Tensor: return _triton_qsa_paged_sparse_gqa( q=q, @@ -298,8 +342,19 @@ def __call__( softmax_scale=softmax_scale, query_positions=query_positions, compress_ratio=compress_ratio, + output_gate=output_gate, ) + @staticmethod + def fuses_output_gate( + q: torch.Tensor, + k_cache: torch.Tensor, + selected_tokens: torch.Tensor, + output_gate: torch.Tensor, + ) -> bool: + """Whether a call with these operands and no prefill pair carries `output_gate`.""" + return _can_fuse_qsa_splitk_output_gate(q, k_cache, selected_tokens, output_gate) + def reference( self, *, @@ -313,8 +368,10 @@ def reference( softmax_scale: float, query_positions: Optional[torch.Tensor] = None, compress_ratio: Optional[int] = None, + output_gate: Optional[torch.Tensor] = None, ) -> torch.Tensor: - """fp32 sparse GQA over every resident selected token, rounded once. + """fp32 sparse GQA over every resident selected token, rounded once; + with `output_gate`, then gated in fp32 and rounded again. `query_positions` and `compress_ratio` are ignored: on a prefix-packed selection the prefill kernel's bound reads exactly the valid entries. @@ -354,7 +411,10 @@ def reference( weights = torch.nan_to_num(torch.softmax(scores, dim=-1), nan=0.0) pooled = torch.einsum("bhgk,bkhd->bhgd", weights, values) out[start:stop] = pooled.reshape(stop - start, num_q_heads, head_dim) - return out.to(q.dtype) + if output_gate is None: + return out.to(q.dtype) + gate = output_gate.reshape(rows, num_q_heads, head_dim).float() + return (out.to(q.dtype).float() * torch.sigmoid(gate)).to(q.dtype) qsa_paged_sparse_gqa = QsaPagedSparseGqa() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_decode.py new file mode 100644 index 000000000000..33274e1c2b2b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_decode.py @@ -0,0 +1,828 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""QSA sparse GQA for decode rows: a split-K pass that loads its tokens in one round trip, +then a merge spread over the SMs. + +Tier 5 (CuTe DSL written for this target), both launches. + +Split-K pass. One CTA (eight warps) covers `CAP` consecutive selected tokens of one KV +head of one row, so a decode row's 2051 selected tokens take 33 CTAs a KV head: + +1. before its PDL wait the CTA stages the KV head's query heads in shared memory and from + there in every warp's MMA A fragments, stages the request's block-table row (up to + `BT_STAGE` pages) in shared memory, and releases its dependents; +2. after it, its tokens' pages come from shared memory and every token's K and V row + (16 bytes a lane, a warp eight tokens) is requested at once with cp.async, so the CTA + pays one DRAM round trip -- QSA's split-K kernel walks two or three 16-token tiles, + each a chain of dependent loads (selection, page, K/V); +3. scores = q K^T on warp MMAs (bf16 in, fp32 out; a warp eight tokens), scaled into the + log2 domain and masked; one softmax over the CTA's tokens (a warp two heads); P V on + warp MMAs with P in bf16 (a warp 32 head dims); each head's partial times its + reciprocal sum is staged in shared memory and stored in 16-byte pieces with its + log2-sum-exp, as QSA's split-K pass stores them. + +Merge. A CTA covers 128 elements of a head (two CTAs a head) with up to eight warps: warp +`w` reads its splits' partials of four adjacent elements a lane, one 16-byte load a split, +all in flight; every warp forms the split weights itself from two LSEs a lane; the warps' +shares are added in shared memory. + +Both launches release their dependents at entry, before their own waits: the merge and, +through it, the output projection that follows launch while the top-k ahead of the +split-K pass still runs, and the projection streams its weight through that latency-bound +stretch. Before its wait the split-K pass reads q and the merge the output gate: the +caller's contract is that both are complete when the grid ahead of the split-K pass +releases it (the target's top-k does so only after its own wait; a launch without PDL +releases at its end). +""" + +from __future__ import annotations + +import functools +import math +import os +from typing import Optional + +import cuda.bindings.driver as cuda_drv +import cutlass +import cutlass.cute as cute +import torch +from cutlass import const_expr +from cutlass.cute import experimental as cute_ext +from cutlass.cute.nvgpu import cpasync, warp +from cutlass.cute.runtime import from_dlpack, make_ptr + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +# Approved exception (as for `qsa_paged_sparse_gqa`): QSA's sparse path has no torch.ops +# surface; its row-count PDL policy is reached as Python. +from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import _qsa_pdl_enabled + +from .qsa_paged_sparse_gqa import qsa_paged_sparse_gqa as _composition + +#: Output elements a merge CTA covers (four adjacent ones a lane of each warp). +CHUNK = 128 +#: Splits the merge reads: two LSEs a lane. +MAX_SPLITS = 64 +#: Rows the target calls this entry with: the launches that run with PDL (QSA enables it up +#: to four rows); larger batches keep the composition. +MAX_ROWS = 4 +#: Selected tokens a split-K CTA covers. +CAP = 64 +#: MMA rows: the query heads of a KV head, padded. +M = 16 +#: bf16 elements of padding after each shared K, V, q and P row (keeps the MMA fragment +#: reads free of bank conflicts). +PAD = 8 +WARPS = 8 +THREADS = WARPS * 32 +#: Block-table entries staged in shared memory before the wait; a wider table is read +#: from global memory after it. +BT_STAGE = 4096 +HEAD_DIM = 256 +LOG2E = 1.4426950408889634 + + +@functools.lru_cache(maxsize=None) +def pdl_enabled(device_index: int) -> bool: + """Programmatic dependent launch: sm_90+ unless TRTLLM_ENABLE_PDL=0.""" + props = torch.cuda.get_device_properties(device_index) + return props.major >= 9 and os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" + + +class _SplitKKernel: + def __init__( + self, + *, + group, + width, + splits, + tpb, + scale, + bt_width, + num_pages, + k_strides, + v_strides, + num_requests, + pdl, + ): + assert group <= M and CAP == 2 * 32 and CAP % WARPS == 0 + self.group, self.width, self.splits, self.tpb = group, width, splits, tpb + self.scale = scale * LOG2E + self.bt_width, self.num_pages, self.num_requests = bt_width, num_pages, num_requests + self.stage_bt = bt_width <= BT_STAGE + self.k_strides, self.v_strides = k_strides, v_strides + self.ld = HEAD_DIM + PAD + self.pdl = pdl + + def smem_bytes(self) -> int: + def r(n): + return -(-n // 128) * 128 + + b = ( + r(M * self.ld * 2) + + 2 * r(CAP * self.ld * 2) + + r(M * (CAP + PAD) * 2) + + r(M * (CAP + 4) * 4) + ) + b += r(M * 2 * 4) + 2 * r(CAP * 8) + if self.stage_bt: + b += r(self.bt_width * 4) + return b + 256 + + @cute.jit + def __call__(self, mQ, pK, pV, mBT, mSel, mReq, pPart, mLse, stream: cuda_drv.CUstream): + rows = cute.size(mSel, mode=[0]) + kv_heads = cute.size(mQ, mode=[1]) // self.group + mma = cute.make_tiled_mma( + warp.MmaF16BF16Op(cutlass.BFloat16, cutlass.Float32, (16, 8, 16)), + cute.make_layout((1, WARPS, 1)), + ) + g2s = cute.make_copy_atom(cpasync.CopyG2SOp(), cutlass.BFloat16, num_bits_per_copy=128) + st = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=128 + ) + self.kernel(mQ, pK, pV, mBT, mSel, mReq, pPart, mLse, mma, g2s, st).launch( + grid=[rows * kv_heads * self.splits, 1, 1], + block=[THREADS, 1, 1], + smem=self.smem_bytes(), + stream=stream, + use_pdl=self.pdl, + ) + + @cute.kernel + def kernel( + self, + mQ, + pK, + pV, + mBT, + mSel, + mReq, + pPart, + mLse, + mma: cute.TiledMma, + g2s: cute.CopyAtom, + st: cute.CopyAtom, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + warp_i = cute.arch.warp_idx() + lane = cute.arch.lane_idx() + group: cutlass.Constexpr = self.group + splits: cutlass.Constexpr = self.splits + width: cutlass.Constexpr = self.width + tpb: cutlass.Constexpr = self.tpb + bt_width: cutlass.Constexpr = self.bt_width + ld: cutlass.Constexpr = self.ld + dim: cutlass.Constexpr = HEAD_DIM + heads = cute.size(mQ, mode=[1]) + kv_heads = heads // group + rows = cute.size(mSel, mode=[0]) + row = bidx // (kv_heads * splits) + kvh = (bidx // splits) % kv_heads + split = bidx % splits + c0 = split * CAP + + smem = cutlass.utils.SmemAllocator() + sQ = smem.allocate_tensor( + cutlass.BFloat16, cute.make_layout((M, dim), stride=(ld, 1)), byte_alignment=128 + ) + sK = smem.allocate_tensor( + cutlass.BFloat16, cute.make_layout((CAP, dim), stride=(ld, 1)), byte_alignment=128 + ) + sV = smem.allocate_tensor( + cutlass.BFloat16, cute.make_layout((CAP, dim), stride=(ld, 1)), byte_alignment=128 + ) + sP = smem.allocate_tensor( + cutlass.BFloat16, cute.make_layout((M, CAP), stride=(CAP + PAD, 1)), byte_alignment=128 + ) + sS = smem.allocate_tensor( + cutlass.Float32, cute.make_layout((M, CAP), stride=(CAP + 4, 1)), byte_alignment=128 + ) + sStat = smem.allocate_tensor( + cutlass.Float32, cute.make_layout((M, 2), stride=(2, 1)), byte_alignment=128 + ) + sKo = smem.allocate_tensor( + cutlass.Int64, cute.make_layout((CAP,), stride=(1,)), byte_alignment=128 + ) + sVo = smem.allocate_tensor( + cutlass.Int64, cute.make_layout((CAP,), stride=(1,)), byte_alignment=128 + ) + + # 1. Before the wait: the KV head's query heads (complete before the grid ahead of + # this one released it) into shared memory and the A fragments, the block-table row + # (host-filled), the release. + cpr: cutlass.Constexpr = dim // 8 # 16-byte pieces a row: one a lane + q_vec = cute.logical_divide(mQ, (None, None, 8)) + sq_vec = cute.logical_divide(sQ, (None, 8)) + for i in cutlass.range_constexpr(-(-(group * cpr) // THREADS)): + c = tidx + i * THREADS + cc = cutlass.min(c, group * cpr - 1) + h = cc // cpr + ch = cc % cpr + if c < group * cpr: + cute.copy(g2s, q_vec[row, kvh * group + h, (None, ch)], sq_vec[h, (None, ch)]) + cute.arch.cp_async_commit_group() + for i in cutlass.range_constexpr((M - group) * dim // THREADS): + e = tidx + i * THREADS + sQ[group + e // dim, e % dim] = cutlass.BFloat16(0.0) + req = mReq[row] + req_ok = (req >= 0) & (req < self.num_requests) + safe_req = cutlass.min(cutlass.max(req, 0), self.num_requests - 1) + if const_expr(self.stage_bt): + sBT = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((bt_width,), stride=(1,)), byte_alignment=128 + ) + for i in cutlass.range_constexpr(-(-bt_width // THREADS)): + e = tidx + i * THREADS + ec = cutlass.min(e, bt_width - 1) + v = mBT[safe_req, ec] + if e < bt_width: + sBT[e] = v + cute.arch.cp_async_wait_group(0) + cute.arch.sync_threads() + thr = mma.get_slice(tidx) + tCsQ = thr.partition_A(sQ) + tCrQ = thr.make_fragment_A(tCsQ) + cute.autovec_copy(tCsQ, tCrQ) + if const_expr(self.pdl): + # The merge waits before reading the partials and reads nothing before its wait + # that the grid ahead of this one writes. + cute.arch.griddepcontrol_launch_dependents() + cute.arch.griddepcontrol_wait() + + # 2. The CTA's tokens: each one's K and V row offsets in elements, or -1. + ks0, ks1, ks2, _ = self.k_strides + vs0, vs1, vs2, _ = self.v_strides + if tidx < CAP: + col = c0 + tidx + colc = cutlass.min(col, width - 1) + tok = mSel[row, colc] + ok = req_ok & (col < width) & (tok >= 0) + safe_tok = cutlass.max(tok, 0) + lp = safe_tok // tpb + ok = ok & (lp < bt_width) + lpc = cutlass.min(lp, bt_width - 1) + page = cutlass.Int32(0) + if const_expr(self.stage_bt): + page = sBT[lpc] + else: + page = mBT[safe_req, lpc] + ok = ok & (page >= 0) & (page < self.num_pages) + pg = cutlass.Int64(cutlass.max(page, 0)) + tip = cutlass.Int64(safe_tok % tpb) + hk = cutlass.Int64(kvh) + sKo[tidx] = cutlass.Int64( + cutlass.select_(ok, pg * ks0 + hk * ks1 + tip * ks2, cutlass.Int64(-1)) + ) + sVo[tidx] = cutlass.Int64( + cutlass.select_(ok, pg * vs0 + hk * vs1 + tip * vs2, cutlass.Int64(-1)) + ) + cute.arch.sync_threads() + + # Every token's K and V row: warp w the tokens 8w .. 8w + 7, a lane 16 bytes of each; + # an invalid token's rows are zeros (its score is masked; V must stay finite). + sk_vec = cute.logical_divide(sK, (None, 8)) + sv_vec = cute.logical_divide(sV, (None, 8)) + k_base = pK.toint() + cutlass.Int64(lane * 16) + v_base = pV.toint() + cutlass.Int64(lane * 16) + tpw: cutlass.Constexpr = CAP // WARPS + for i in cutlass.range_constexpr(tpw): + t = warp_i * tpw + i + ko = sKo[t] + vo = sVo[t] + ksrc = cute.make_tensor( + cute.make_ptr( + cutlass.BFloat16, + k_base + cutlass.max(ko, 0) * 2, + cute.AddressSpace.gmem, + assumed_align=16, + ), + cute.make_layout((8,), stride=(1,)), + ) + vsrc = cute.make_tensor( + cute.make_ptr( + cutlass.BFloat16, + v_base + cutlass.max(vo, 0) * 2, + cute.AddressSpace.gmem, + assumed_align=16, + ), + cute.make_layout((8,), stride=(1,)), + ) + if ko >= 0: + cute.copy(g2s, ksrc, sk_vec[t, (None, lane)]) + cute.copy(g2s, vsrc, sv_vec[t, (None, lane)]) + else: + for e in cutlass.range_constexpr(8): + sK[t, lane * 8 + e] = cutlass.BFloat16(0.0) + sV[t, lane * 8 + e] = cutlass.BFloat16(0.0) + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) + cute.arch.sync_threads() + + # 3. Scores: a warp eight tokens over the whole head dim, masked into shared memory. + tCsK = thr.partition_B(sK) + tCrK = thr.make_fragment_B(tCsK) + cute.autovec_copy(tCsK, tCrK) + scores = cute.make_rmem_tensor(thr.partition_shape_C((M, CAP)), cutlass.Float32) + scores.fill(0.0) + cute.gemm(mma, scores, tCrQ, tCrK, scores) + cs = thr.partition_C(cute.make_identity_tensor((M, CAP))) + neg_inf = cutlass.Float32(-math.inf) + for j in cutlass.range_constexpr(cute.size(scores)): + h, t = cs[j] + v = neg_inf + if sKo[t] >= 0: + v = scores[j] * self.scale + sS[h, t] = v + cute.arch.sync_threads() + # The softmax: warp w the heads w and w + WARPS, a lane two tokens; padding heads + # get P = 0. + for j in cutlass.range_constexpr(M // WARPS): + h = warp_i + WARPS * j + a = sS[h, lane] + b = sS[h, lane + 32] + mx = cute.arch.warp_reduction_max(cute.arch.fmax(a, b)) + safe_mx = cutlass.Float32(cutlass.select_(mx == neg_inf, cutlass.Float32(0.0), mx)) + keep = h < group + pa = cutlass.Float32( + cutlass.select_(keep, cute.math.exp2(a - safe_mx), cutlass.Float32(0.0)) + ) + pb = cutlass.Float32( + cutlass.select_(keep, cute.math.exp2(b - safe_mx), cutlass.Float32(0.0)) + ) + den = cute.arch.warp_reduction_sum(pa + pb) + sP[h, lane] = pa.to(cutlass.BFloat16) + sP[h, lane + 32] = pb.to(cutlass.BFloat16) + if lane == 0: + sStat[h, 0] = mx + sStat[h, 1] = den + cute.arch.sync_threads() + + # P V: a warp 32 of the head dims (four 8-wide tiles). + tCsP = thr.partition_A(sP) + tCrP = thr.make_fragment_A(tCsP) + cute.autovec_copy(tCsP, tCrP) + sVt = cute.make_tensor(sV.iterator, cute.make_layout((dim, CAP), stride=(1, ld))) + tCsV = thr.partition_B(sVt) + tCrV = thr.make_fragment_B(tCsV) + cute.autovec_copy(tCsV, tCrV) + acc = cute.make_rmem_tensor(thr.partition_shape_C((M, dim)), cutlass.Float32) + acc.fill(0.0) + cute.gemm(mma, acc, tCrP, tCrV, acc) + + # Each head's partial times its reciprocal sum, staged over sK (free since the + # softmax barrier) as [group, dim] fp32, then stored in 16-byte pieces: the CTA's + # heads are contiguous in the partial tensor. + sOut = cute.make_tensor( + cute.recast_ptr(sK.iterator, dtype=cutlass.Float32), + cute.make_layout((group, dim), stride=(dim, 1)), + ) + co = thr.partition_C(cute.make_identity_tensor((M, dim))) + h0 = lane // 4 # an m16n8 C fragment's rows: h0 and h0 + 8 + d0 = sStat[h0, 1] + d1 = sStat[h0 + 8, 1] + r0 = cutlass.Float32(0.0) + r1 = cutlass.Float32(0.0) + if d0 > 0.0: + r0 = 1.0 / d0 + if d1 > 0.0: + r1 = 1.0 / d1 + for j in cutlass.range_constexpr(cute.size(acc)): + h, d = co[j] + r = cutlass.Float32(cutlass.select_(h < 8, r0, r1)) + if h < group: + sOut[h, d] = acc[j] * r + cute.arch.sync_threads() + base = ((cutlass.Int64(split) * rows + row) * heads + kvh * group) * dim + out_ptr = pPart.toint() + base * 4 + so_vec = cute.make_tensor( + sOut.iterator, cute.make_layout((group * dim // 4, 4), stride=(4, 1)) + ) + reg = cute.make_rmem_tensor(cute.make_layout((4,), stride=(1,)), cutlass.Float32) + for i in cutlass.range_constexpr(-(-(group * dim // 4) // THREADS)): + c = tidx + i * THREADS + if c < group * dim // 4: + cute.autovec_copy(so_vec[c, None], reg) + dst = cute.make_tensor( + cute.make_ptr( + cutlass.Float32, + out_ptr + cutlass.Int64(c * 16), + cute.AddressSpace.gmem, + assumed_align=16, + ), + cute.make_layout((4,), stride=(1,)), + ) + cute.copy(st, reg, dst) + if tidx < group: + mx = sStat[tidx, 0] + den = sStat[tidx, 1] + lse = neg_inf + if den > 0.0: + lse = mx + cute.math.log2(den) + mLse[split, row, kvh * group + tidx] = lse + + +class _MergeKernel: + def __init__(self, *, heads, dim, splits, gated, pdl): + assert dim % CHUNK == 0 and 1 < splits <= MAX_SPLITS + self.heads, self.dim, self.splits, self.gated, self.pdl = heads, dim, splits, gated, pdl + self.chunks = dim // CHUNK + self.warps = min(8, splits) + self.spw = -(-splits // self.warps) # splits a warp, rounded up + + @cute.jit + def __call__(self, mPart, mLse, mGate, mOut, stream: cuda_drv.CUstream): + rows = cute.size(mOut, mode=[0]) + ld = cute.make_copy_atom(cute.nvgpu.CopyG2ROp(), cutlass.Float32, num_bits_per_copy=128) + self.kernel(mPart, mLse, mGate, mOut, ld).launch( + grid=[rows * self.heads * self.chunks, 1, 1], + block=[self.warps * 32, 1, 1], + smem=self.warps * CHUNK * 4 + 64, + stream=stream, + use_pdl=self.pdl, + ) + + @cute.kernel + def kernel(self, mPart, mLse, mGate, mOut, ld: cute.CopyAtom): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + warp_i = cute.arch.warp_idx() + lane = cute.arch.lane_idx() + per_row: cutlass.Constexpr = self.heads * self.chunks + spw: cutlass.Constexpr = self.spw + row = bidx // per_row + head = (bidx % per_row) // self.chunks + c0 = (bidx % self.chunks) * CHUNK + red = cutlass.utils.SmemAllocator().allocate_tensor( + cutlass.Float32, + cute.make_layout((self.warps, CHUNK), stride=(CHUNK, 1)), + byte_alignment=16, + ) + # The gate, before the wait: see the module notes. + g = cutlass.Float32(0.0) + if const_expr(self.gated): + if tidx < CHUNK: + g = mGate[row, head * self.dim + c0 + tidx].to(cutlass.Float32) + if const_expr(self.pdl): + # The output projection after this reads only its weight before its own wait. + cute.arch.griddepcontrol_launch_dependents() + cute.arch.griddepcontrol_wait() + + # 1. The split weights, in every warp: exp2(lse - max), and their sum. + neg_inf = cutlass.Float32(-math.inf) + l0 = neg_inf + l1 = neg_inf + if lane < self.splits: + l0 = mLse[lane, row, head] + if lane + 32 < self.splits: + l1 = mLse[lane + 32, row, head] + # 2. Warp w's splits w * spw .. w * spw + spw - 1, four adjacent elements a lane: + # one 16-byte load a split, all in flight. A slot past the split count reads the + # last split again and weighs it 0. + part_vec = cute.logical_divide(mPart, (None, None, None, 4)) + p = cute.make_rmem_tensor(cute.make_layout((spw, 4), stride=(4, 1)), cutlass.Float32) + for j in cutlass.range_constexpr(spw): + sidx = cutlass.min(warp_i * spw + j, self.splits - 1) + cute.copy(ld, part_vec[sidx, row, head, (None, (c0 + 4 * lane) // 4)], p[j, None]) + mx = cute.arch.warp_reduction_max(cute.arch.fmax(l0, l1)) + w0 = cutlass.Float32(0.0) + w1 = cutlass.Float32(0.0) + if mx > neg_inf: + w0 = cute.math.exp2(l0 - mx) + w1 = cute.math.exp2(l1 - mx) + den = cute.arch.warp_reduction_sum(w0 + w1) + acc = cute.make_rmem_tensor(cute.make_layout((4,), stride=(1,)), cutlass.Float32) + acc.fill(0.0) + for j in cutlass.range_constexpr(spw): + sidx = warp_i * spw + j + wa = cute.arch.shuffle_sync(w0, sidx % 32) + wb = cute.arch.shuffle_sync(w1, sidx % 32) + w = cutlass.Float32(cutlass.select_(sidx < 32, wa, wb)) + w = cutlass.Float32(cutlass.select_(sidx < self.splits, w, cutlass.Float32(0.0))) + for e in cutlass.range_constexpr(4): + acc[e] = w * p[j, e] + acc[e] + for e in cutlass.range_constexpr(4): + red[warp_i, 4 * lane + e] = acc[e] + cute.arch.sync_threads() + + # 3. The warps' shares summed in warp order; the output gate on the rounded + # attention value, as the composition applies it. + if tidx < CHUNK: + tot = cutlass.Float32(0.0) + for w_ in cutlass.range_constexpr(self.warps): + tot = tot + red[w_, tidx] + merged = cutlass.Float32(0.0) + if den > 0.0: + merged = tot / den + if const_expr(self.gated): + att = merged.to(cutlass.BFloat16).to(cutlass.Float32) + merged = att * (1.0 / (1.0 + cute.math.exp(-g))) + mOut[row, head, c0 + tidx] = merged.to(mOut.element_type) + + +def _dl(t: torch.Tensor, align: int = 4): + return from_dlpack(t, assumed_align=align) + + +def _ptr(t: torch.Tensor, dtype): + return make_ptr(dtype, t.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) + + +def splits(q, k_cache, selected_tokens) -> int: + """The split-K CTAs a row and KV head: one per `CAP` selected tokens.""" + return -(-selected_tokens.shape[1] // CAP) + + +@functools.lru_cache(maxsize=None) +def _compiled_splitk( + rows, + q_shape, + q_stride, + group, + width, + tpb, + scale, + bt_shape, + num_pages, + k_strides, + v_strides, + pdl, +): + heads = q_shape[1] + n_splits = -(-width // CAP) + kern = _SplitKKernel( + group=group, + width=width, + splits=n_splits, + tpb=tpb, + scale=scale, + bt_width=bt_shape[1], + num_pages=num_pages, + k_strides=k_strides, + v_strides=v_strides, + num_requests=bt_shape[0], + pdl=pdl, + ) + dev = "cuda" + q = torch.empty_strided(q_shape, q_stride, dtype=torch.bfloat16, device=dev) + kv = torch.empty(64, dtype=torch.bfloat16, device=dev) + bt = torch.empty(bt_shape, dtype=torch.int32, device=dev) + sel = torch.empty(rows, width, dtype=torch.int32, device=dev) + req = torch.empty(rows, dtype=torch.int32, device=dev) + part = torch.empty(64, dtype=torch.float32, device=dev) + lse = torch.empty(n_splits, rows, heads, dtype=torch.float32, device=dev) + stream = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream) + return cute_ext.compile( + kern, + _dl(q, 16), + _ptr(kv, cutlass.BFloat16), + _ptr(kv, cutlass.BFloat16), + _dl(bt), + _dl(sel), + _dl(req), + _ptr(part, cutlass.Float32), + _dl(lse), + stream, + ) + + +@functools.lru_cache(maxsize=None) +def _compiled_merge(rows, heads, dim, n_splits, gate_stride, dtype, pdl): + kern = _MergeKernel( + heads=heads, dim=dim, splits=n_splits, gated=gate_stride is not None, pdl=pdl + ) + dev = "cuda" + part = torch.empty(n_splits, rows, heads, dim, dtype=torch.float32, device=dev) + lse = torch.empty(n_splits, rows, heads, dtype=torch.float32, device=dev) + gate = torch.empty(rows, gate_stride or heads * dim, dtype=dtype, device=dev)[:, : heads * dim] + out = torch.empty(rows, heads, dim, dtype=dtype, device=dev) + stream = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream) + return cute_ext.compile(kern, _dl(part, 16), _dl(lse), _dl(gate, 2), _dl(out, 2), stream) + + +def launch( + q, + k_cache, + v_cache, + block_table, + selected_tokens, + request_indices, + tokens_per_block, + softmax_scale, + output_gate=None, + pdl=True, +): + """The launch behind `QsaPagedSparseGqaDecode.__call__`; returns `[rows, H, D]` in q's dtype.""" + rows, num_q_heads, head_dim = q.shape + num_kv_heads = k_cache.shape[1] + group_size = num_q_heads // num_kv_heads + width = selected_tokens.shape[1] + n_splits = -(-width // CAP) + output = torch.empty((rows, num_q_heads, head_dim), dtype=q.dtype, device=q.device) + if rows == 0: + return output + dev = q.device + pdl = pdl and pdl_enabled(dev.index if dev.index is not None else torch.cuda.current_device()) + use_pdl = bool(pdl and _qsa_pdl_enabled(rows)) + partial_output = torch.empty( + (n_splits, rows, num_q_heads, head_dim), dtype=torch.float32, device=dev + ) + partial_lse = torch.empty((n_splits, rows, num_q_heads), dtype=torch.float32, device=dev) + fn = _compiled_splitk( + rows, + tuple(q.shape), + tuple(q.stride()), + group_size, + width, + tokens_per_block, + float(softmax_scale), + tuple(block_table.shape), + k_cache.shape[0], + tuple(k_cache.stride()), + tuple(v_cache.stride()), + use_pdl, + ) + stream = cuda_drv.CUstream(torch.cuda.current_stream(dev).cuda_stream) + fn( + _dl(q, 16), + _ptr(k_cache, cutlass.BFloat16), + _ptr(v_cache, cutlass.BFloat16), + _dl(block_table), + _dl(selected_tokens), + _dl(request_indices), + _ptr(partial_output, cutlass.Float32), + _dl(partial_lse), + stream, + ) + gate2d = None if output_gate is None else output_gate.reshape(rows, num_q_heads * head_dim) + mfn = _compiled_merge( + rows, + num_q_heads, + head_dim, + n_splits, + None if gate2d is None else gate2d.stride(0), + q.dtype, + use_pdl, + ) + mfn( + _dl(partial_output, 16), + _dl(partial_lse), + _dl(output.view(rows, num_q_heads * head_dim) if gate2d is None else gate2d, 2), + _dl(output, 2), + stream, + ) + return output + + +def supports(q, k_cache, selected_tokens) -> bool: + """Whether these operands are ones this entry serves (see `is_valid`).""" + return ( + q.dim() == 3 + and 0 < q.shape[0] <= MAX_ROWS + and q.shape[2] == HEAD_DIM + and q.dtype == torch.bfloat16 + and k_cache.dtype == torch.bfloat16 + and q.shape[1] % k_cache.shape[1] == 0 + and q.shape[1] // k_cache.shape[1] <= M + and 1 < splits(q, k_cache, selected_tokens) <= MAX_SPLITS + and all(t.data_ptr() % 16 == 0 for t in (q, k_cache)) + and all(s % 8 == 0 for s in (*q.stride()[:2], *k_cache.stride()[:3])) + ) + + +class QsaPagedSparseGqaDecode(OpWrapper): + """`qsa_paged_sparse_gqa` without a prefill pair, for decode rows: a split-K pass whose + CTAs load their tokens' K/V in one round trip, then the merge -- and, with + `output_gate`, the gate -- as one launch. Returns a new `[rows, H, D]` tensor in q's + dtype. + + Fusion boundary: the split-K attention, the merge, the optional sigmoid output gate. + The selection, the K/V store and the output projection are the caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="decode, 1 row at position 8700, gated", spec=dict(pos=(8700,), gated=True)), + Cell(why="decode, 1 row at position 8700, no gate", spec=dict(pos=(8700,), gated=False)), + Cell(why="decode, 2 rows, gated", spec=dict(pos=(8700, 5000), gated=True)), + Cell( + why="decode, 4 rows, one a short sequence, gated", + spec=dict(pos=(8700, 5000, 77, 3000), gated=True), + ), + Cell( + why="decode, 1 row with half the selection -1", + spec=dict(pos=(8700,), gated=True, holes=True), + ), + ) + + note = """ + Metadata consumed: as `qsa_paged_sparse_gqa` (block table, request indices). Preconditions: + as that entry's, plus: no prefill pair; at most `MAX_ROWS` rows; bf16 q and K/V with a head + dim of 256 and at most 16 query heads a KV head; more than one and at most 64 split-K CTAs + a row and KV head (`splits(...)`: 64 selected tokens a CTA); `output_gate` (when given) + `[rows, H * D]` or `[rows, H, D]` with unit inner stride in q's dtype. `supports(...)` is + the predicate a caller can check. + + PDL: the split-K pass reads q and the block table before its wait and the merge reads the + output gate before its own; the selection and the K/V rows are read after the split-K + pass's wait. So q and the gate must be complete when the grid ahead of this entry releases + it: a launch without PDL releases at its end; a PDL producer of the selection (the + target's top-k) must release only after its own wait. + + Numerics: one softmax over a CTA's 64 tokens (QSA's kernel runs an online softmax over + 16-token tiles) and the merge sums the splits in split order (QSA's merge reduces them as + a tree), so the fp32 values can differ in the last bits and the bf16 output by an ulp on + some elements. + """ + + def __call__( + self, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + block_table: torch.Tensor, + selected_tokens: torch.Tensor, + request_indices: torch.Tensor, + tokens_per_block: int, + softmax_scale: float, + output_gate: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return launch( + q, + k_cache, + v_cache, + block_table, + selected_tokens, + request_indices, + tokens_per_block, + softmax_scale, + output_gate, + ) + + def reference( + self, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + block_table: torch.Tensor, + selected_tokens: torch.Tensor, + request_indices: torch.Tensor, + tokens_per_block: int, + softmax_scale: float, + output_gate: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """The composition entry's reference with the same operands.""" + return _composition.reference( + q=q, + k_cache=k_cache, + v_cache=v_cache, + block_table=block_table, + selected_tokens=selected_tokens, + request_indices=request_indices, + tokens_per_block=tokens_per_block, + softmax_scale=softmax_scale, + query_positions=None, + compress_ratio=None, + output_gate=output_gate, + ) + + def is_valid( + self, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + block_table: torch.Tensor, + selected_tokens: torch.Tensor, + request_indices: torch.Tensor, + tokens_per_block: int, + softmax_scale: float, + output_gate: Optional[torch.Tensor] = None, + ) -> None: + assert q.ndim == 3 and q.shape[2] == HEAD_DIM, "q [rows, H, 256]" + assert 0 < q.shape[0] <= MAX_ROWS, f"1 .. {MAX_ROWS} rows" + assert q.dtype == torch.bfloat16 and k_cache.dtype == v_cache.dtype == torch.bfloat16, ( + "bf16 q and K/V" + ) + assert q.shape[1] % k_cache.shape[1] == 0 and q.shape[1] // k_cache.shape[1] <= M, ( + "<= 16 q heads a KV head" + ) + assert 1 < splits(q, k_cache, selected_tokens) <= MAX_SPLITS, ( + "65 .. 4096 selected tokens a row" + ) + assert q.stride(2) == 1 and k_cache.stride(3) == 1 and v_cache.stride(3) == 1, ( + "unit inner strides" + ) + assert all(t.data_ptr() % 16 == 0 for t in (q, k_cache, v_cache)), ( + "16-byte aligned q and K/V" + ) + assert all( + s % 8 == 0 for s in (*q.stride()[:2], *k_cache.stride()[:3], *v_cache.stride()[:3]) + ), "q and K/V rows 16-byte aligned" + if output_gate is not None: + assert output_gate.dtype == q.dtype and output_gate.stride(-1) == 1, "gate in q's dtype" + assert output_gate.numel() == q.numel(), "gate covers q" + + +qsa_paged_sparse_gqa_decode = QsaPagedSparseGqaDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_rowtile_attention_cute.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_rowtile_attention_cute.py new file mode 100644 index 000000000000..3b6b861108d9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_rowtile_attention_cute.py @@ -0,0 +1,998 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The row-tile pass of `qsa_sparse_gqa_rowtile` on SM103 tensor cores (CuTe DSL, tcgen05). + +Same inputs and output as that op's Triton `_rowtile_attention_kernel`: per group of +`rows_per_tile` rows of one request, the plan's packed slot list (`slots[group]`, `listed` +entries), each row's bitmap re-indexed onto it (`packed_bits`, as int32 halves) and +`group_plan[group] = (tiles, listed)`; groups with `tiles < 0` belong to the lone-row pass. + +One CTA per (group, KV head), 128 MMA rows = `rows_per_tile` query rows x `group_size` +heads, 9 warps: + +* warps 5-8 gather the list's K and V rows (64 a tile) with `cp.async` into one ring of + `kv_slots` 32 KB SW128 shared-memory slots, filled in the MMA warp's order K0, K1, V0, + K2, V1, ...; each slot completes on its mbarrier through `cp.async.mbarrier.arrive`; +* warp 4 issues the tcgen05 MMAs: S[j % 2] = Q K(j)^T (Q from tensor memory), then + O += P(j - 1) V(j - 1) (P from tensor memory, overwriting the S it came from); +* warps 0-3, one thread per MMA row, first copy their Q row into tensor memory, then per + tile load S with its row max (`tcgen05.ld.red`), rescale O when that max grew by more + than `rescale_threshold` in log2 units (lazily; otherwise exponentiate against the stale + max), write P = 2^(s - max) on the row's own bits and 0 elsewhere, and at the end + normalize O and store it through the idle ring as whole 512-byte rows. + +Tensor memory: O in columns 0-255, S/P buffers at 256 and 320, Q at 384 (512 columns). +The tile's row max covers every listed entry, an upper bound of the row's own: a valid +softmax offset. P is rounded to bf16 for the PV MMA and summed in fp32 unrounded, as the +Triton kernel does. +""" + +from __future__ import annotations + +import functools + +import cuda.bindings.driver as cuda_drv +import cutlass +import cutlass.cute as cute +import cutlass.cute.nvgpu.tcgen05 as tcgen05 +import cutlass.pipeline as pipeline +import cutlass.utils.blackwell_helpers as sm100_utils +import torch +from cutlass._mlir.dialects import nvvm as _nvvm +from cutlass.cute.nvgpu import OperandMajorMode, cpasync +from cutlass.cute.runtime import make_ptr +from cutlass.cute.typing import Float32, Int32, Int64 +from cutlass.cutlass_dsl import dsl_user_op + +try: + from cutlass.memory import SmemAllocator +except ImportError: # DSL releases before the memory namespace split + from cutlass.utils import SmemAllocator + +BF16 = cutlass.BFloat16 +_LOG2_E = 1.4426950408889634 +#: The head dimension the tensor-memory plan is laid out for. +HEAD_DIM = 256 +#: MMA rows per CTA and tokens per tile (one 64-bit bitmap word). +BLOCK_M = 128 +BLOCK_N = 64 + + +@dsl_user_op +def tcgen05_fence_before(*, loc=None, ip=None): + """tcgen05.fence::before_thread_sync: order this thread's prior tcgen05 ops before a thread sync.""" + _nvvm.tcgen05_fence(_nvvm.Tcgen05FenceKind.BEFORE_THREAD_SYNC, loc=loc, ip=ip) + + +@dsl_user_op +def tcgen05_fence_after(*, loc=None, ip=None): + """tcgen05.fence::after_thread_sync: order later tcgen05 ops after a thread sync.""" + _nvvm.tcgen05_fence(_nvvm.Tcgen05FenceKind.AFTER_THREAD_SYNC, loc=loc, ip=ip) + + +class _RowtileAttentionKernel: + def __init__( + self, + group_size: int, + rows_per_tile: int, + head_dim: int = 256, + tokens_per_block: int = 32, + kv_slots: int = 6, + rescale_threshold: float = 8.0, + use_pdl: bool = True, + ): + assert head_dim == 256, ( + "the TMEM plan (O 256 + 2 x S 64 + Q 128 columns) is laid out for D = 256" + ) + assert group_size * rows_per_tile <= 128 + self.group_size = group_size + self.rows_per_tile = rows_per_tile + self.d = head_dim + self.tpb = tokens_per_block + self.bm = 128 + self.bn = 64 + self.kv_slots = kv_slots + self.tau = rescale_threshold + self.use_pdl = use_pdl + self.mma_warp = 4 + self.load_warp0 = 5 + self.threads = 32 * 9 + self.tmem_o = 0 + self.tmem_s0 = 256 + self.tmem_q = 384 # Q as the QK MMA's A operand: 256 bf16 a row, two a column + self.tmem_cols = 512 + self.chunk_cols = 32 # O columns a TMEM load/store moves per thread + + @cute.jit + def __call__( + self, + q: cute.Pointer, + k: cute.Pointer, + v: cute.Pointer, + slots: cute.Pointer, + bits: cute.Pointer, + plan: cute.Pointer, + out: cute.Pointer, + rows: Int32, + num_groups: Int32, + num_kv_heads: Int32, + q_stride_row: Int32, + k_stride_page: Int32, + k_stride_head: Int32, + k_stride_token: Int32, + v_stride_page: Int32, + v_stride_head: Int32, + v_stride_token: Int32, + slots_stride_group: Int32, + bits_stride_row: Int32, + out_stride_row: Int32, + scale_log2: Float32, + stream, + ): + def qk(): + return sm100_utils.make_trivial_tiled_mma( + BF16, + BF16, + OperandMajorMode.K, + OperandMajorMode.K, + Float32, + tcgen05.CtaGroup.ONE, + (self.bm, self.bn), + tcgen05.OperandSource.TMEM, + ) + + def pv(): + return sm100_utils.make_trivial_tiled_mma( + BF16, + BF16, + OperandMajorMode.K, + OperandMajorMode.MN, + Float32, + tcgen05.CtaGroup.ONE, + (self.bm, self.d), + tcgen05.OperandSource.TMEM, + ) + + # Overwrite / accumulate twins of each MMA: the ACCUMULATE field is set once at kernel + # scope, so no loop or branch rebinds an MMA value. + qk_mma, qk_mma_acc, pv_mma, pv_mma_acc = qk(), qk(), pv(), pv() + q_layout = sm100_utils.make_smem_layout_a(qk_mma, (self.bm, self.bn, self.d), BF16, 1) + k_layout = sm100_utils.make_smem_layout_b( + qk_mma, (self.bm, self.bn, self.d), BF16, self.kv_slots + ) + v_layout = sm100_utils.make_smem_layout_b( + pv_mma, (self.bm, self.d, self.bn), BF16, self.kv_slots + ) + p_layout = sm100_utils.make_smem_layout_a(pv_mma, (self.bm, self.d, self.bn), BF16, 1) + + @cute.struct + class SharedStorage: + q_full: cute.struct.MemRange[Int64, 1] + kv_full: cute.struct.MemRange[Int64, self.kv_slots] + kv_empty: cute.struct.MemRange[Int64, self.kv_slots] + s_full: cute.struct.MemRange[Int64, 2] + p_full: cute.struct.MemRange[Int64, 2] + o_full: cute.struct.MemRange[Int64, 1] + o_final: cute.struct.MemRange[Int64, 1] + tmem_holding: Int32 + sKV: cute.struct.Align[cute.struct.MemRange[BF16, cute.cosize(k_layout)], 1024] + + self.shared_storage = SharedStorage + + self.kernel( + qk_mma, + qk_mma_acc, + pv_mma, + pv_mma_acc, + q_layout, + k_layout, + v_layout, + p_layout, + q, + k, + v, + slots, + bits, + plan, + out, + rows, + q_stride_row, + k_stride_page, + k_stride_head, + k_stride_token, + v_stride_page, + v_stride_head, + v_stride_token, + slots_stride_group, + bits_stride_row, + out_stride_row, + scale_log2, + ).launch( + grid=[num_groups, num_kv_heads, 1], + block=[self.threads, 1, 1], + stream=stream, + min_blocks_per_mp=1, + use_pdl=self.use_pdl, + ) + + @cute.kernel + def kernel( + self, + qk_mma: cute.TiledMma, + qk_mma_acc: cute.TiledMma, + pv_mma: cute.TiledMma, + pv_mma_acc: cute.TiledMma, + q_layout: cute.ComposedLayout, + k_layout: cute.ComposedLayout, + v_layout: cute.ComposedLayout, + p_layout: cute.ComposedLayout, + q: cute.Pointer, + k: cute.Pointer, + v: cute.Pointer, + slots: cute.Pointer, + bits: cute.Pointer, + plan: cute.Pointer, + out: cute.Pointer, + rows: Int32, + q_stride_row: Int32, + k_stride_page: Int32, + k_stride_head: Int32, + k_stride_token: Int32, + v_stride_page: Int32, + v_stride_head: Int32, + v_stride_token: Int32, + slots_stride_group: Int32, + bits_stride_row: Int32, + out_stride_row: Int32, + scale_log2: Float32, + ): + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tidx, _, _ = cute.arch.thread_idx() + group, kv_head, _ = cute.arch.block_idx() + qk_mma.set(tcgen05.Field.ACCUMULATE, False) + qk_mma_acc.set(tcgen05.Field.ACCUMULATE, True) + pv_mma.set(tcgen05.Field.ACCUMULATE, False) + pv_mma_acc.set(tcgen05.Field.ACCUMULATE, True) + + mPlan = cute.make_tensor(plan, cute.make_layout((2 * 1048576,))) + tiles = mPlan[group * 2] + listed = mPlan[group * 2 + 1] + + if tiles >= 0: + smem = SmemAllocator() + storage = smem.allocate(self.shared_storage) + q_full = storage.q_full.data_ptr() + kv_full = storage.kv_full.data_ptr() + kv_empty = storage.kv_empty.data_ptr() + s_full = storage.s_full.data_ptr() + p_full = storage.p_full.data_ptr() + o_full = storage.o_full.data_ptr() + o_final = storage.o_final.data_ptr() + if tidx == 0: + cute.arch.mbarrier_init(q_full, 128) + for i in cutlass.range_constexpr(self.kv_slots): + cute.arch.mbarrier_init(kv_full + i, 128) + cute.arch.mbarrier_init(kv_empty + i, 1) + for i in cutlass.range_constexpr(2): + cute.arch.mbarrier_init(s_full + i, 1) + cute.arch.mbarrier_init(p_full + i, 128) + cute.arch.mbarrier_init(o_full, 1) + cute.arch.mbarrier_init(o_final, 1) + cute.arch.mbarrier_init_fence() + + tmem_users = pipeline.NamedBarrier(barrier_id=1, num_threads=32 * 5) + if warp_idx == self.mma_warp: + cute.arch.alloc_tmem(self.tmem_cols, storage.tmem_holding.ptr) + cute.arch.relinquish_tmem_alloc_permit() + cute.arch.sync_threads() + + # Slices taken before the role branches: the MMA warp's `set(ACCUMULATE)` rebinds the + # tiled MMAs inside its branch, which the other branches must not see. + qk_thr = qk_mma.get_slice(0) + pv_thr = pv_mma.get_slice(0) + sK = storage.sKV.get_tensor(k_layout.outer, swizzle=k_layout.inner) + sV = cute.make_tensor(cute.recast_ptr(sK.iterator, v_layout.inner), v_layout.outer) + + if warp_idx >= self.load_warp0: + self.load_role( + tidx - 32 * self.load_warp0, + group, + kv_head, + tiles, + listed, + sK, + k, + v, + slots, + k_stride_page, + k_stride_head, + k_stride_token, + v_stride_page, + v_stride_head, + v_stride_token, + slots_stride_group, + kv_full, + kv_empty, + ) + elif warp_idx == self.mma_warp: + tmem_ptr = cute.arch.retrieve_tmem_ptr(Float32, 16, storage.tmem_holding.ptr) + self.mma_role( + qk_mma, + qk_mma_acc, + pv_mma, + pv_mma_acc, + qk_thr, + pv_thr, + q_layout, + p_layout, + sK, + sV, + tmem_ptr, + tiles, + q_full, + kv_full, + kv_empty, + s_full, + p_full, + o_full, + o_final, + ) + tmem_users.arrive_and_wait() + cute.arch.dealloc_tmem(tmem_ptr, self.tmem_cols) + else: + tmem_ptr = cute.arch.retrieve_tmem_ptr(Float32, 16, storage.tmem_holding.ptr) + self.softmax_role( + qk_thr, + pv_thr, + tmem_ptr, + tidx, + group, + kv_head, + tiles, + rows, + bits, + bits_stride_row, + out, + out_stride_row, + scale_log2, + q_full, + s_full, + p_full, + o_full, + o_final, + storage.sKV.data_ptr(), + q, + q_stride_row, + ) + tmem_users.arrive_and_wait() + + if cutlass.const_expr(self.use_pdl): + # The lone-row pass may still run: finish only after it has, so a later launch sees every row. + cute.arch.griddepcontrol_wait() + + # ------------------------------------------------------------------ loaders (warps 5-8) + @cute.jit + def load_role( + self, + lt: Int32, + group: Int32, + kv_head: Int32, + tiles: Int32, + listed: Int32, + sK, + k: cute.Pointer, + v: cute.Pointer, + slots: cute.Pointer, + k_stride_page: Int32, + k_stride_head: Int32, + k_stride_token: Int32, + v_stride_page: Int32, + v_stride_head: Int32, + v_stride_token: Int32, + slots_stride_group: Int32, + kv_full, + kv_empty, + ): + """Fill the K/V ring in the MMA warp's order K0, K1, V0, K2, V1, ..., V(T-1). Thread `lt` + gathers tokens 16 (lt // 32) + 4 t + (lt % 32) // 8 (t < 4), one 16-byte chunk (lt % 8) of + each of their four 128-byte dim blocks; the slots of tile j + 1 are fetched while tile j's + copies issue.""" + D = self.d + atom = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), BF16, num_bits_per_copy=128 + ) + chunk = lt % 8 + # Element-coordinate view of the swizzled ring: (8 elements, (chunk, 64-dim block), token, slot). + sKV_c = cute.make_tensor( + sK.iterator, + cute.make_layout( + (8, (8, D // 64), self.bn, self.kv_slots), + stride=(1, (8, self.bn * 64), 64, self.bn * D), + ), + ) + if tiles > 0: + mSlots = cute.make_tensor( + slots + Int64(group) * slots_stride_group, cute.make_layout((1 << 30,)) + ) + first_token = 16 * (lt // 32) + (lt % 32) // 8 + slot_next = cute.make_rmem_tensor((4,), Int32) + k_cur = cute.make_rmem_tensor((4,), Int64) + v_cur = cute.make_rmem_tensor((4,), Int64) + v_prev = cute.make_rmem_tensor((4,), Int64) + for t in cutlass.range_constexpr(4): + slot_next[t] = mSlots[cutlass.min(first_token + 4 * t, listed - 1)] + ring = Int32(0) + ring_phase = Int32(1) + for j in cutlass.range(0, tiles, 1, unroll=1): + for t in cutlass.range_constexpr(4): + slot = slot_next[t] + page = slot // self.tpb + tok = slot % self.tpb + k_cur[t] = ( + Int64(page) * k_stride_page + kv_head * k_stride_head + tok * k_stride_token + ) + v_cur[t] = ( + Int64(page) * v_stride_page + kv_head * v_stride_head + tok * v_stride_token + ) + if j + 1 < tiles: + for t in cutlass.range_constexpr(4): + slot_next[t] = mSlots[ + cutlass.min((j + 1) * self.bn + first_token + 4 * t, listed - 1) + ] + # K(j) + cute.arch.mbarrier_wait(kv_empty + ring, ring_phase) + self.issue_tile(atom, k, k_cur, sKV_c, ring, chunk, first_token) + cute.arch.cp_async_mbarrier_arrive_noinc(kv_full + ring) + ring, ring_phase = self.advance(ring, ring_phase) + if j > 0: + # V(j - 1) + cute.arch.mbarrier_wait(kv_empty + ring, ring_phase) + self.issue_tile(atom, v, v_prev, sKV_c, ring, chunk, first_token) + cute.arch.cp_async_mbarrier_arrive_noinc(kv_full + ring) + ring, ring_phase = self.advance(ring, ring_phase) + for t in cutlass.range_constexpr(4): + v_prev[t] = v_cur[t] + # V(tiles - 1) + cute.arch.mbarrier_wait(kv_empty + ring, ring_phase) + self.issue_tile(atom, v, v_prev, sKV_c, ring, chunk, first_token) + cute.arch.cp_async_mbarrier_arrive_noinc(kv_full + ring) + ring, ring_phase = self.advance(ring, ring_phase) + # Tail: every slot released, so no load is in flight when the CTA retires and the ring + # is free for the epilogue's staging. + for i in cutlass.range_constexpr(self.kv_slots): + cute.arch.mbarrier_wait(kv_empty + ring, ring_phase) + ring, ring_phase = self.advance(ring, ring_phase) + + @cute.jit + def issue_tile( + self, + atom, + src_base: cute.Pointer, + bases, + sKV_c, + ring: Int32, + chunk: Int32, + first_token: Int32, + ): + for t in cutlass.range_constexpr(4): + n = first_token + 4 * t + for b in cutlass.range_constexpr(self.d // 64): + offset = cute.assume(bases[t] + b * 64 + chunk * 8, divby=8) + src = cute.make_tensor(src_base + offset, cute.make_layout((8,))) + dst = sKV_c[(None, (chunk, b), n, ring)] + cute.copy(atom, src, dst) + + # ------------------------------------------------------------------ MMA (warp 4) + @cute.jit + def mma_role( + self, + qk_mma: cute.TiledMma, + qk_mma_acc: cute.TiledMma, + pv_mma: cute.TiledMma, + pv_mma_acc: cute.TiledMma, + qk_thr, + pv_thr, + q_layout, + p_layout, + sK, + sV, + tmem_ptr, + tiles: Int32, + q_full, + kv_full, + kv_empty, + s_full, + p_full, + o_full, + o_final, + ): + """QK(0), then QK(j) / PV(j - 1) pairs, then the last PV. S/P buffer b sits at TMEM columns + tmem_s0 + 64 b and is picked by pointer offset; QK(j) overwrites S[j % 2] only after PV(j - 2) + was issued (tcgen05 MMAs of one thread execute in order), and P(j - 2) was its last reader.""" + tQ = cute.make_tensor(tmem_ptr + self.tmem_q, q_layout.outer) + tSrQ = qk_thr.make_fragment_A(tQ)[None, None, None, 0] + tSrQ = cute.make_tensor(cute.recast_ptr(tmem_ptr + self.tmem_q, dtype=BF16), tSrQ.layout) + tSrK = qk_thr.make_fragment_B(sK) + tOrV = pv_thr.make_fragment_B(sV) + tStS = qk_thr.make_fragment_C(qk_thr.partition_shape_C((self.bm, self.bn))) + tOtO_f = pv_thr.make_fragment_C(pv_thr.partition_shape_C((self.bm, self.d))) + tOtO = cute.make_tensor(tmem_ptr + self.tmem_o, tOtO_f.layout) + tP = cute.make_tensor(tmem_ptr + self.tmem_s0, p_layout.outer) + tOrP = pv_thr.make_fragment_A(tP)[None, None, None, 0] + if tiles > 0: + cute.arch.mbarrier_wait(q_full, 0) + tcgen05_fence_after() + ring = Int32(0) + ring_phase = Int32(0) + # QK(0) + cute.arch.mbarrier_wait(kv_full + ring, ring_phase) + self.gemm_qk( + qk_mma, + qk_mma_acc, + tSrQ, + tSrK[None, None, None, ring], + cute.make_tensor(tmem_ptr + self.tmem_s0, tStS.layout), + ) + with cute.arch.elect_one(): + tcgen05.commit(s_full) + tcgen05.commit(kv_empty + ring) + ring, ring_phase = self.advance(ring, ring_phase) + for j in cutlass.range(1, tiles + 1, 1, unroll=1): + buf = j % 2 + if j < tiles: + # QK(j) into S[j % 2] + cute.arch.mbarrier_wait(kv_full + ring, ring_phase) + self.gemm_qk( + qk_mma, + qk_mma_acc, + tSrQ, + tSrK[None, None, None, ring], + cute.make_tensor(tmem_ptr + (self.tmem_s0 + buf * self.bn), tStS.layout), + ) + with cute.arch.elect_one(): + tcgen05.commit(s_full + buf) + tcgen05.commit(kv_empty + ring) + ring, ring_phase = self.advance(ring, ring_phase) + # PV(j - 1) + prev = j - 1 + pbuf = prev % 2 + cute.arch.mbarrier_wait(kv_full + ring, ring_phase) + cute.arch.mbarrier_wait(p_full + pbuf, (prev // 2) % 2) + tcgen05_fence_after() + tOrP_b = cute.make_tensor( + cute.recast_ptr( + tmem_ptr + (self.tmem_s0 + pbuf * self.bn), dtype=tOrP.element_type + ), + tOrP.layout, + ) + if prev == 0: + self.gemm_pv(pv_mma, pv_mma_acc, tOrP_b, tOrV[None, None, None, ring], tOtO) + else: + self.gemm_pv(pv_mma_acc, pv_mma_acc, tOrP_b, tOrV[None, None, None, ring], tOtO) + with cute.arch.elect_one(): + tcgen05.commit(kv_empty + ring) + tcgen05.commit(o_full) + ring, ring_phase = self.advance(ring, ring_phase) + with cute.arch.elect_one(): + tcgen05.commit(o_final) + + @cute.jit + def advance(self, ring: Int32, ring_phase: Int32): + ring = ring + 1 + if ring == self.kv_slots: + ring = Int32(0) + ring_phase = ring_phase ^ 1 + return ring, ring_phase + + @cute.jit + def gemm_qk(self, first: cute.TiledMma, rest: cute.TiledMma, tSrQ, tSrK, tStS): + for kb in cutlass.range_constexpr(cute.size(tSrQ, mode=[2])): + if cutlass.const_expr(kb == 0): + cute.gemm(first, tStS, tSrQ[None, None, kb], tSrK[None, None, kb], tStS) + else: + cute.gemm(rest, tStS, tSrQ[None, None, kb], tSrK[None, None, kb], tStS) + + @cute.jit + def gemm_pv(self, first: cute.TiledMma, rest: cute.TiledMma, tOrP, tOrV, tOtO): + for kb in cutlass.range_constexpr(cute.size(tOrP, mode=[2])): + if cutlass.const_expr(kb == 0): + cute.gemm(first, tOtO, tOrP[None, None, kb], tOrV[None, None, kb], tOtO) + else: + cute.gemm(rest, tOtO, tOrP[None, None, kb], tOrV[None, None, kb], tOtO) + + # ------------------------------------------------------------------ softmax + epilogue (warps 0-3) + @cute.jit + def softmax_role( + self, + qk_thr, + pv_thr, + tmem_ptr, + st: Int32, + group: Int32, + kv_head: Int32, + tiles: Int32, + rows: Int32, + bits: cute.Pointer, + bits_stride_row: Int32, + out: cute.Pointer, + out_stride_row: Int32, + scale_log2: Float32, + q_full, + s_full, + p_full, + o_full, + o_final, + stage, + q: cute.Pointer, + q_stride_row: Int32, + ): + G = self.group_size + R = self.rows_per_tile + D = self.d + tStS = qk_thr.make_fragment_C(qk_thr.partition_shape_C((self.bm, self.bn))) + tStS0 = cute.make_tensor(tmem_ptr + self.tmem_s0, tStS.layout) + tOtO_f = pv_thr.make_fragment_C(pv_thr.partition_shape_C((self.bm, self.d))) + tOtO = cute.make_tensor(tmem_ptr + self.tmem_o, tOtO_f.layout) + + # S: one 64-column load a tile; P: 32 packed-bf16 columns stored over the S it came from. + cS = cute.make_identity_tensor((self.bm, self.bn)) + tScS = qk_thr.partition_C(cS) + # The load also reduces the row max over the tile's 64 columns (tcgen05.ld.red, SM103): an + # upper bound of the row's masked max, which serves as the softmax offset just as well. + load_s = tcgen05.make_tmem_copy( + cute.make_copy_atom( + tcgen05.copy.LdRed32x32bOp( + tcgen05.copy.Repetition(32), redOp=tcgen05.TmemLoadRedOp.MAX + ), + Float32, + ), + tStS0, + ) + thr_load_s = load_s.get_slice(st) + tLtS0 = thr_load_s.partition_S(tStS0) + tLcS = thr_load_s.partition_D(tScS) + p_cols = self.bn // 2 + tStP_layout = cute.composition(tStS0.layout, cute.make_layout((self.bm, p_cols))) + tStP0 = cute.make_tensor(tStS0.iterator, tStP_layout) + tScP = cute.make_tensor( + tScS.iterator, cute.composition(tScS.layout, cute.make_layout((self.bm, p_cols))) + ) + store_p = tcgen05.make_tmem_copy( + cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(p_cols)), Float32), + tStP0, + ) + thr_store_p = store_p.get_slice(st) + tPtP0 = thr_store_p.partition_D(tStP0) + tPcP = thr_store_p.partition_S(tScP) + + # O in 32-column chunks (rescale and epilogue). + cO = cute.make_identity_tensor((self.bm, self.d)) + tOcO = pv_thr.partition_C(cO) + tOtO_c = cute.make_tensor( + tOtO.iterator, + cute.composition(tOtO.layout, cute.make_layout((self.bm, self.chunk_cols))), + ) + tOcO_c = cute.make_tensor( + tOcO.iterator, + cute.composition(tOcO.layout, cute.make_layout((self.bm, self.chunk_cols))), + ) + load_o = tcgen05.make_tmem_copy( + cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(self.chunk_cols)), Float32 + ), + tOtO_c, + ) + store_o = tcgen05.make_tmem_copy( + cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(self.chunk_cols)), Float32 + ), + tOtO_c, + ) + thr_load_o = load_o.get_slice(st) + thr_store_o = store_o.get_slice(st) + tLtO = thr_load_o.partition_S(tOtO_c) + tLcO = thr_load_o.partition_D(tOcO_c) + tStO = thr_store_o.partition_D(tOtO_c) + n_chunks = self.d // self.chunk_cols + + m = st + local_row = m // G + head = m % G + row = group * R + local_row + live = (local_row < R) and (row < rows) + safe_row = group * R + if live: + safe_row = row + mBits = cute.make_tensor( + bits + Int64(safe_row) * bits_stride_row, cute.make_layout((1 << 30,)) + ) + + # Q row -> TMEM columns tmem_q.. (the QK MMA's A operand), 64 elements a store. Dead rows + # store zeros; their S rows are masked out. + if tiles > 0: + q_row = cute.make_tensor( + q + cute.assume(Int64(safe_row) * q_stride_row + (kv_head * G + head) * D, divby=8), + cute.make_layout((D,)), + ) + for c in cutlass.range_constexpr(D // (2 * p_cols)): + rQ = cute.make_rmem_tensor(tPcP.shape, Float32) + rQ16 = cute.make_tensor( + cute.recast_ptr(rQ.iterator, dtype=BF16), cute.make_layout((2 * p_cols,)) + ) + if live: + for i in cutlass.range_constexpr(2 * p_cols // 8): + cute.autovec_copy( + cute.make_tensor( + q_row.iterator + (c * 2 * p_cols + i * 8), cute.make_layout((8,)) + ), + cute.make_tensor(rQ16.iterator + i * 8, cute.make_layout((8,))), + ) + else: + rQ.fill(0.0) + cute.copy( + store_p, + rQ, + cute.make_tensor( + tPtP0.iterator + (self.tmem_q - self.tmem_s0 + c * p_cols), tPtP0.layout + ), + ) + cute.arch.fence_view_async_tmem_store() + tcgen05_fence_before() + cute.arch.mbarrier_arrive(q_full) + + row_max = -Float32.inf + row_sum = Float32(0.0) + lo_next = Int32(0) + hi_next = Int32(0) + if live and tiles > 0: + lo_next = mBits[0] + hi_next = mBits[1] + tau = Float32(self.tau) + for j in cutlass.range(0, tiles, 1, unroll=1): + lo = lo_next + hi = hi_next + if live and j + 1 < tiles: + lo_next = mBits[2 * j + 2] + hi_next = mBits[2 * j + 3] + buf = j % 2 + cute.arch.mbarrier_wait(s_full + buf, (j // 2) % 2) + tcgen05_fence_after() + rS = cute.make_rmem_tensor(tLcS.shape, Float32) + n_ld = cute.size(rS, mode=[1]) + rMax = cute.make_rmem_tensor(cute.make_layout((1, n_ld)), Float32) + tLtS_b = cute.make_tensor(tLtS0.iterator + buf * self.bn, tLtS0.layout) + for c in cutlass.range_constexpr(n_ld): + cute.copy_atom_call( + load_s, tLtS_b[None, c, 0, 0], (rS[None, c, 0, 0], rMax[None, c]) + ) + cute.arch.fence_view_async_tmem_load() + tile_max = rMax[0] + for c in cutlass.range_constexpr(1, n_ld): + tile_max = cute.arch.fmax(tile_max, rMax[c]) + new_max = cute.arch.fmax(row_max, tile_max) + use_max = row_max + factor = Float32(1.0) + if row_max == -Float32.inf: + use_max = new_max + elif (new_max - row_max) * scale_log2 > tau: + use_max = new_max + factor = cute.math.exp2((row_max - new_max) * scale_log2, fastmath=True) + need = factor != Float32(1.0) + if cute.arch.vote_any_sync(need): + # O holds PV(0..j-1): wait for PV(j-1), then rescale this warp's rows. + cute.arch.mbarrier_wait(o_full, (j - 1) % 2) + tcgen05_fence_after() + for c in cutlass.range_constexpr(n_chunks): + rO = cute.make_rmem_tensor(tLcO.shape, Float32) + cute.copy( + load_o, + cute.make_tensor(tLtO.iterator + c * self.chunk_cols, tLtO.layout), + rO, + ) + for i in cutlass.range_constexpr(cute.size(rO)): + rO[i] = rO[i] * factor + cute.copy( + store_o, + rO, + cute.make_tensor(tStO.iterator + c * self.chunk_cols, tStO.layout), + ) + cute.arch.fence_view_async_tmem_store() + row_max = use_max + neg = -use_max * scale_log2 + rP = cute.make_rmem_tensor(tPcP.shape, Float32) + rP16 = cute.make_tensor(cute.recast_ptr(rP.iterator, dtype=BF16), rS.layout) + # p = 2^(s * scale - max * scale) on the row's own entries, 0 elsewhere; four partial sums. + acc0 = (Float32(0.0), Float32(0.0)) + acc1 = (Float32(0.0), Float32(0.0)) + for i in cutlass.range_constexpr(0, self.bn, 4): + word = lo if i < 32 else hi + e = cute.make_rmem_tensor((4,), Float32) + e[0], e[1] = cute.arch.fma_packed_f32x2( + (rS[i], rS[i + 1]), (scale_log2, scale_log2), (neg, neg) + ) + e[2], e[3] = cute.arch.fma_packed_f32x2( + (rS[i + 2], rS[i + 3]), (scale_log2, scale_log2), (neg, neg) + ) + for u in cutlass.range_constexpr(4): + x = cute.math.exp2(e[u], fastmath=True) + rS[i + u] = x if (word & (1 << ((i + u) % 32))) != 0 else Float32(0.0) + acc0 = cute.arch.add_packed_f32x2(acc0, (rS[i], rS[i + 1])) + acc1 = cute.arch.add_packed_f32x2(acc1, (rS[i + 2], rS[i + 3])) + row_sum = row_sum * factor + ((acc0[0] + acc0[1]) + (acc1[0] + acc1[1])) + rP16.store(rS.load().to(BF16)) + cute.copy(store_p, rP, cute.make_tensor(tPtP0.iterator + buf * self.bn, tPtP0.layout)) + cute.arch.fence_view_async_tmem_store() + tcgen05_fence_before() + cute.arch.mbarrier_arrive(p_full + buf) + + # Epilogue: O / l in bf16, staged through the (now idle) K/V ring -- a row's 32 16-byte + # chunks XOR-swizzled by the row -- so each warp then stores whole 512-byte rows. + if tiles > 0: + cute.arch.mbarrier_wait(o_final, 0) + tcgen05_fence_after() + inv = Float32(0.0) + if row_sum > Float32(0.0): + inv = Float32(1.0) / row_sum + swz = m % 32 + for c in cutlass.range_constexpr(n_chunks): + rO = cute.make_rmem_tensor(tLcO.shape, Float32) + if tiles > 0: + cute.copy( + load_o, cute.make_tensor(tLtO.iterator + c * self.chunk_cols, tLtO.layout), rO + ) + for i in cutlass.range_constexpr(cute.size(rO)): + rO[i] = rO[i] * inv + else: + rO.fill(0.0) + rB = cute.make_rmem_tensor((self.chunk_cols,), BF16) + rB.store(rO.load().to(BF16)) + for i in cutlass.range_constexpr(self.chunk_cols // 8): + logical = c * (self.chunk_cols // 8) + i + offset = cute.assume(m * D + ((logical ^ swz) * 8), divby=8) + dst = cute.make_tensor(stage + offset, cute.make_layout((8,))) + cute.autovec_copy( + cute.make_tensor(rB.iterator + i * 8, cute.make_layout((8,))), dst + ) + cute.arch.sync_warp() + lane = st % 32 + warp = st // 32 + for r in cutlass.range_constexpr(32): + mr = warp * 32 + r + lr = mr // G + row_r = group * R + lr + if lr < R and row_r < rows: + src_off = cute.assume(mr * D + ((lane ^ r) * 8), divby=8) + rV = cute.make_rmem_tensor((8,), BF16) + cute.autovec_copy(cute.make_tensor(stage + src_off, cute.make_layout((8,))), rV) + dst_off = cute.assume( + Int64(row_r) * out_stride_row + (kv_head * G + mr % G) * D + lane * 8, divby=8 + ) + cute.autovec_copy(rV, cute.make_tensor(out + dst_off, cute.make_layout((8,)))) + + +@functools.lru_cache(maxsize=None) +def _compiled(group_size: int, rows_per_tile: int, tokens_per_block: int, use_pdl: bool): + kernel = _RowtileAttentionKernel( + group_size, rows_per_tile, HEAD_DIM, tokens_per_block, use_pdl=use_pdl + ) + gmem = cute.AddressSpace.gmem + + def fake(dtype, align): + return make_ptr(dtype, 16 * 1024, gmem, assumed_align=align) + + args = ( + fake(BF16, 16), + fake(BF16, 16), + fake(BF16, 16), + fake(Int32, 4), + fake(Int32, 4), + fake(Int32, 4), + fake(BF16, 16), + 1, + 1, + 1, + 8, + 8, + 8, + 8, + 8, + 8, + 8, + 1, + 1, + 8, + 1.0, + cuda_drv.CUstream(0), + ) + return cute.compile(kernel, *args) + + +def _aligned(t: torch.Tensor, strides: int) -> bool: + return t.data_ptr() % 16 == 0 and all(t.stride(i) % 8 == 0 for i in range(strides)) + + +def eligible( + q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, block_m: int, block_n: int +) -> bool: + """Whether a launch fits this pass: SM103, bf16 q and K/V, D = 256, the 128-row tile of + 64-token tiles, unit-stride head dimension and 16-byte-aligned rows.""" + if (block_m, block_n) != (BLOCK_M, BLOCK_N) or q.shape[2] != HEAD_DIM: + return False + if not (q.dtype == k_cache.dtype == v_cache.dtype == torch.bfloat16): + return False + if torch.cuda.get_device_capability(q.device) != (10, 3): + return False + if ( + q.stride(2) != 1 + or q.stride(1) != HEAD_DIM + or k_cache.stride(3) != 1 + or v_cache.stride(3) != 1 + ): + return False + return _aligned(q, 1) and _aligned(k_cache, 3) and _aligned(v_cache, 3) + + +class _Launcher: + """`rowtile_attention_cute[grid](...)`: grid-indexed like the Triton kernel it stands in for; + the grid is `(groups, kv_heads)`.""" + + def __getitem__(self, grid): + return functools.partial(self._launch, tuple(grid)) + + @staticmethod + def _launch( + grid, + q, + k_cache, + v_cache, + slots, + packed_bits, + group_plan, + output, + softmax_scale, + rows, + *, + group_size, + rows_per_tile, + tokens_per_block, + use_pdl, + ): + assert grid == (group_plan.shape[0], k_cache.shape[1]), grid + assert ( + output.is_contiguous() + and packed_bits.dtype == torch.int64 + and slots.dtype == torch.int32 + ) + fn = _compiled(group_size, rows_per_tile, tokens_per_block, bool(use_pdl)) + gmem = cute.AddressSpace.gmem + + def ptr(t, dtype, align): + return make_ptr(dtype, t.data_ptr(), gmem, assumed_align=align) + + fn( + ptr(q, BF16, 16), + ptr(k_cache, BF16, 16), + ptr(v_cache, BF16, 16), + ptr(slots, Int32, 4), + ptr(packed_bits, Int32, 4), + ptr(group_plan, Int32, 4), + ptr(output, BF16, 16), + rows, + grid[0], + grid[1], + q.stride(0), + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(2), + v_cache.stride(0), + v_cache.stride(1), + v_cache.stride(2), + slots.stride(0), + 2 * packed_bits.stride(0), + output.stride(0), + softmax_scale * _LOG2_E, + cuda_drv.CUstream(torch.cuda.current_stream(q.device).cuda_stream), + ) + + +rowtile_attention_cute = _Launcher() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py new file mode 100644 index 000000000000..b26d04ebb77f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py @@ -0,0 +1,1043 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""QSA sparse GQA for context batches: several query rows per program, K/V read once per tile. + +The first launch turns each row's selection (the columns the prefill bound +lets it read) into a bitmap over logical token positions. A plan launch forms, +per group of consecutive rows of one request, the union of its rows' resident +selected tokens as a packed list of physical cache slots, and each row's +bitmap re-indexed onto that list. The row-tile launch runs one program per +(group, KV head) that walks the packed list in 64-slot tiles, each tile's K/V +loaded once for every row and head of the group, each row masked to its own +bits. A group whose rows span requests (generation rows of a mixed batch) +walks each row's own selected columns instead, as the per-row kernel does. +""" + +from __future__ import annotations + +import math +import os + +import cutlass +import cutlass.cute as cute +import torch +import triton +import triton.language as tl +from cutlass._mlir.dialects import llvm +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.cute_dsl_kernels.blackwell.utils import make_ptr + +from .qsa_paged_sparse_gqa import qsa_paged_sparse_gqa as _per_row_entry +from .qsa_rowtile_attention_cute import eligible as _cute_eligible +from .qsa_rowtile_attention_cute import rowtile_attention_cute as _rowtile_attention_cute + +_LOG2_E = 1.4426950408889634 +_INTEGER_DTYPES = (torch.int32, torch.int64) +# The per-row kernel's scan tile: the prefill bound is rounded up to it, so the +# bitmap reads exactly the columns that kernel reads. +_BOUND_TILE = 64 +# Programs of the lone-row pass per KV head; each walks the lone list with this stride. +_LONE_PROGRAMS = 256 # capacity: fixed: a launch width; the lone list may be any length +# Bitmap words the plan pass unions and packs per loop step. +_PLAN_WORDS = 16 # capacity: fixed: a loop step; the plan loops over every word + + +def _pdl_enabled() -> bool: + if os.environ.get("TRTLLM_ENABLE_PDL", "1") != "1": + return False + return torch.cuda.get_device_capability()[0] >= 9 + + +@triton.jit +def _rowtile_bits_kernel( + selected_tokens, + query_positions, + bits, + row_info, + selected_stride_row, + selected_stride_token, + bits_stride_row, + TOPK: tl.constexpr, + BLOCK_TOPK: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + TOKENS_PER_BLOCK: tl.constexpr, + PAGE_TABLE_WIDTH: tl.constexpr, + BOUND_TILE: tl.constexpr, + BLOCK_C: tl.constexpr, +): + """Row `r`: OR a bit per readable selected token into `bits[r]`; record the + highest such token and the number of columns read in `row_info[r]`.""" + row = tl.program_id(0).to(tl.int64) + visible_tokens = tl.load(query_positions + row) + 1 + selected_blocks = tl.minimum(visible_tokens // COMPRESS_RATIO, BLOCK_TOPK) + valid_tokens = selected_blocks * COMPRESS_RATIO + visible_tokens % COMPRESS_RATIO + columns_read = tl.minimum( + tl.cdiv(tl.minimum(valid_tokens, TOPK), BOUND_TILE) * BOUND_TILE, TOPK + ) + columns_read = tl.maximum(columns_read, 0).to(tl.int32) + offsets = tl.arange(0, BLOCK_C) + highest = tl.full([BLOCK_C], -1, tl.int32) + one = tl.full([BLOCK_C], 1, tl.int64) + for start in tl.range(0, columns_read, BLOCK_C): + columns = start + offsets + in_bound = columns < columns_read + tokens = tl.load( + selected_tokens + row * selected_stride_row + columns * selected_stride_token, + mask=in_bound, + other=-1, + ).to(tl.int32) + readable = in_bound & (tokens >= 0) & (tokens // TOKENS_PER_BLOCK < PAGE_TABLE_WIDTH) + safe = tl.maximum(tokens, 0) + tl.atomic_or( + bits + row * bits_stride_row + (safe >> 6), + one << (safe & 63).to(tl.int64), + mask=readable, + sem="relaxed", + ) + highest = tl.maximum(highest, tl.where(readable, tokens, -1)) + tl.store(row_info + row * 2, tl.max(highest, axis=0)) + tl.store(row_info + row * 2 + 1, columns_read) + + +# -------------------------------------------------------------------------------------- +# The bitmap pass in CuTe DSL (Tier 6). The Triton pass above ORs every readable selected +# token's bit into a zeroed global bitmap with a global `atomic_or` (~2051 a row). Here a +# CTA owns one row: the row's selection is staged in shared memory with coalesced loads, +# each thread walks `_BITS_RUN` consecutive columns, ORs the bits of one 64-token word in a +# register and flushes the word into a shared-memory bitmap with one `red.shared.or.b64` +# when the run leaves it; then the CTA writes every word of the row and `row_info`. The +# result is the Triton pass's, word for word, in any column order (OR is idempotent and +# commutative); the order only decides how few flushes a run makes, and the selection +# arrives in token order. No memset of the bitmap is needed. +# -------------------------------------------------------------------------------------- + +_BITS_THREADS = 128 +#: Consecutive selected columns a thread walks. +_BITS_RUN = 16 + + +@dsl_user_op +def _red_or_shared(addr, val, *, loc=None, ip=None): + llvm.inline_asm( + T.i32(), + [cutlass.Int32(addr).ir_value(loc=loc, ip=ip), cutlass.Int64(val).ir_value(loc=loc, ip=ip)], + "{ red.shared.or.b64 [$1], $2; mov.u32 $0, 0; }", + "=r,r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +class _RowBitsKernel: + """`_rowtile_bits_kernel`'s `bits` and `row_info`, one CTA a row.""" + + def __init__(self, topk, block_topk, compress_ratio, tokens_per_block, table_width, num_words): + self.topk, self.block_topk, self.cr = topk, block_topk, compress_ratio + self.tpb, self.width, self.num_words = tokens_per_block, table_width, num_words + self.staged = topk + topk // _BITS_RUN + 1 + + @cute.jit + def __call__( + self, + sel_t: cute.Tensor, + qp_ptr: cute.Pointer, + bits_ptr: cute.Pointer, + info_ptr: cute.Pointer, + rows: cutlass.Int32, + stream, + ): + qp = cute.make_tensor(qp_ptr, cute.make_layout((rows,), stride=(1,))) + bits = cute.make_tensor( + bits_ptr, cute.make_layout((rows, self.num_words), stride=(self.num_words, 1)) + ) + info = cute.make_tensor(info_ptr, cute.make_layout((rows, 2), stride=(2, 1))) + # `sel_t` arrives as a tensor: the TVM-FFI entry takes its stream from a GPU tensor. + self.kernel(sel_t, qp, bits, info).launch( + grid=[rows, 1, 1], + block=[_BITS_THREADS, 1, 1], + smem=self.num_words * 8 + self.staged * 4 + 128, + stream=stream, + ) + + @cute.kernel + def kernel(self, mSel: cute.Tensor, mQp: cute.Tensor, mBits: cute.Tensor, mInfo: cute.Tensor): + tidx, _, _ = cute.arch.thread_idx() + row, _, _ = cute.arch.block_idx() + smem = cutlass.utils.SmemAllocator() + words = smem.allocate_tensor( + cutlass.Int64, cute.make_layout((self.num_words,), stride=(1,)), byte_alignment=16 + ) + red = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((4,), stride=(1,)), byte_alignment=16 + ) + # The row's selection, one padding slot every `_BITS_RUN` columns (conflict-free runs). + sel_s = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((self.staged,), stride=(1,)), byte_alignment=16 + ) + words_base = cutlass.Int32(words.iterator.toint()) + for i in cutlass.range_constexpr((self.num_words + _BITS_THREADS - 1) // _BITS_THREADS): + w = tidx + i * _BITS_THREADS + if w < self.num_words: + words[w] = cutlass.Int64(0) + + # The columns the per-row kernel reads: the prefill bound rounded up to its tile. + visible = cutlass.Int32(mQp[row]) + 1 + selected_blocks = cutlass.min(visible // self.cr, cutlass.Int32(self.block_topk)) + valid = selected_blocks * self.cr + visible % self.cr + bounded = cutlass.min(valid, cutlass.Int32(self.topk)) + columns_read = cutlass.min( + (bounded + _BOUND_TILE - 1) // _BOUND_TILE * _BOUND_TILE, cutlass.Int32(self.topk) + ) + columns_read = cutlass.max(columns_read, cutlass.Int32(0)) + for i in cutlass.range_constexpr((self.topk + _BITS_THREADS - 1) // _BITS_THREADS): + c = tidx + i * _BITS_THREADS + if c < columns_read: + sel_s[c + c // _BITS_RUN] = cutlass.Int32(mSel[row, c]) + cute.arch.sync_threads() + + highest = cutlass.Int32(-1) + runs = (self.topk + _BITS_THREADS * _BITS_RUN - 1) // (_BITS_THREADS * _BITS_RUN) + for run in cutlass.range_constexpr(runs): + c0 = (run * _BITS_THREADS + tidx) * _BITS_RUN + if c0 < columns_read: + cur_w = cutlass.Int32(-1) + cur_m = cutlass.Int64(0) + for k in cutlass.range_constexpr(_BITS_RUN): + c = c0 + k + t = cutlass.Int32(-1) + if c < columns_read: + t = sel_s[c + c // _BITS_RUN] + if t >= 0: + if t // self.tpb < self.width: + w = t >> 6 + if w != cur_w: + if cur_w >= 0: + _red_or_shared(words_base + cur_w * 8, cur_m) + cur_w = w + cur_m = cutlass.Int64(0) + cur_m = cur_m | (cutlass.Int64(1) << cutlass.Int64(t & 63)) + highest = cutlass.max(highest, t) + if cur_w >= 0: + _red_or_shared(words_base + cur_w * 8, cur_m) + for off in cutlass.range_constexpr(5): + highest = cutlass.max(highest, cute.arch.shuffle_sync_bfly(highest, offset=1 << off)) + if tidx % 32 == 0: + red[tidx // 32] = highest + cute.arch.sync_threads() + if tidx == 0: + mInfo[row, 0] = cutlass.max(cutlass.max(red[0], red[1]), cutlass.max(red[2], red[3])) + mInfo[row, 1] = columns_read + for i in cutlass.range_constexpr((self.num_words + _BITS_THREADS - 1) // _BITS_THREADS): + w = tidx + i * _BITS_THREADS + if w < self.num_words: + mBits[row, w] = words[w] + + +_ROW_BITS: dict = {} + + +def _row_bits( + selected_tokens, query_positions, compress_ratio, tokens_per_block, table_width, num_words +): + """`(bits [rows, num_words] int64, row_info [rows, 2] int32)` of the bitmap pass: the + CuTe kernel when the selection has unit column stride, the Triton pass otherwise.""" + rows, final_topk = selected_tokens.shape + block_topk = (final_topk - compress_ratio + 1) // compress_ratio + device = selected_tokens.device + row_info = torch.empty((rows, 2), dtype=torch.int32, device=device) + query_positions = query_positions.contiguous() + if selected_tokens.stride(1) != 1: + bits = torch.zeros((rows, num_words), dtype=torch.int64, device=device) + _rowtile_bits_kernel[(rows,)]( + selected_tokens, query_positions, bits, row_info, + selected_tokens.stride(0), selected_tokens.stride(1), bits.stride(0), + TOPK=final_topk, BLOCK_TOPK=block_topk, COMPRESS_RATIO=compress_ratio, + TOKENS_PER_BLOCK=tokens_per_block, PAGE_TABLE_WIDTH=table_width, + BOUND_TILE=_BOUND_TILE, BLOCK_C=1024, num_warps=4, + ) # fmt: skip + return bits, row_info + bits = torch.empty((rows, num_words), dtype=torch.int64, device=device) + key = ( + final_topk, compress_ratio, tokens_per_block, table_width, + query_positions.dtype, selected_tokens.dtype, + ) # fmt: skip + args = (selected_tokens, query_positions.data_ptr(), bits.data_ptr(), row_info.data_ptr()) + fn = _ROW_BITS.get(key) + if fn is None: + gm = cute.AddressSpace.gmem + qty = cutlass.Int64 if query_positions.dtype == torch.int64 else cutlass.Int32 + fn = cute.compile( + _RowBitsKernel( + final_topk, block_topk, compress_ratio, tokens_per_block, table_width, num_words + ), + from_dlpack(selected_tokens, assumed_align=4).mark_layout_dynamic(leading_dim=1), + make_ptr(qty, args[1], gm, assumed_align=4), + make_ptr(cutlass.Int64, args[2], gm, assumed_align=16), + make_ptr(cutlass.Int32, args[3], gm, assumed_align=8), + cutlass.Int32(1), + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--opt-level 2 --enable-tvm-ffi", + ) + _ROW_BITS[key] = fn + fn(*args, rows) + return bits, row_info + + +@triton.jit +def _or_combine(a, b): + return a | b + + +@triton.jit +def _compress_bits(x, m): + """Per word, the bits of `x` at the set positions of `m`, packed to the low + end in order (a parallel bit extract; Hacker's Delight `compress`, 64-bit). + `x` and `m` are uint64 and broadcast together.""" + x = x & m + mk = ~m << 1 + for i in tl.static_range(6): + mp = mk ^ (mk << 1) + mp = mp ^ (mp << 2) + mp = mp ^ (mp << 4) + mp = mp ^ (mp << 8) + mp = mp ^ (mp << 16) + mp = mp ^ (mp << 32) + mv = mp & m + m = (m ^ mv) | (mv >> (1 << i)) + t = x & mv + x = (x ^ t) | (t >> (1 << i)) + mk = mk & ~mp + return x + + +@triton.jit +def _rowtile_plan_kernel( + request_indices, row_info, bits, block_table, slots, packed_bits, group_plan, + lone_rows, lone_count, rows, + bits_stride_row, block_table_stride_request, block_table_stride_page, slots_stride_group, + ROWS_PER_TILE: tl.constexpr, NUM_REQUESTS: tl.constexpr, ROWS_POW2: tl.constexpr, + BLOCK_N: tl.constexpr, SHARED_TILE_COST: tl.constexpr, NUM_CACHE_PAGES: tl.constexpr, + PAGE_TABLE_WIDTH: tl.constexpr, TOKENS_PER_BLOCK: tl.constexpr, PLAN_WORDS: tl.constexpr, +): # fmt: skip + """Per group of rows of one request: list the union of its rows' resident + selected tokens, in logical order, as physical cache slots + (`page * tokens_per_block + offset`) in `slots[group]`, and write each + row's bitmap re-indexed onto that list into `packed_bits[row]` (bit `j` + set when the row selects the list's `j`-th token). `group_plan[group]` is + (tiles of the list for the row-tile pass, or -1; list length). A group the + row-tile pass does not take -- rows of several requests, or a list whose + tiles cost more than the rows' own columns -- appends its rows to the lone + list instead.""" + group = tl.program_id(0).to(tl.int64) + offsets = tl.arange(0, ROWS_POW2) + members = group * ROWS_PER_TILE + offsets + live = (offsets < ROWS_PER_TILE) & (members < rows) + requests = tl.load(request_indices + members, mask=live, other=-1).to(tl.int64) + lowest = tl.min(tl.where(live, requests, NUM_REQUESTS), axis=0) + highest_request = tl.max(tl.where(live, requests, -1), axis=0) + highest_token = tl.max(tl.load(row_info + members * 2, mask=live, other=-1), axis=0) + columns = tl.load(row_info + members * 2 + 1, mask=live, other=0) + own_tiles = tl.sum(tl.cdiv(columns, BLOCK_N), axis=0) + one_request = (lowest == highest_request) & (lowest >= 0) & (lowest < NUM_REQUESTS) + words_needed = tl.where(one_request, (highest_token + 64) // 64, 0) + page_base = ( + block_table + + tl.minimum(tl.maximum(lowest, 0), NUM_REQUESTS - 1) * block_table_stride_request + ) + word_offsets = tl.arange(0, PLAN_WORDS) + lanes = tl.arange(0, 64) + lane_bits = tl.full([64], 1, tl.uint64) << lanes.to(tl.uint64) + listed = tl.sum(tl.zeros([ROWS_POW2], tl.int32), axis=0) + for start in range(0, words_needed, PLAN_WORDS): + words = start + word_offsets + in_range = words < words_needed + row_words = tl.load( + bits + members[:, None] * bits_stride_row + words[None, :], + mask=live[:, None] & in_range[None, :], + other=0, + ).to(tl.uint64, bitcast=True) # fmt: skip + selected = tl.reduce(row_words, 0, _or_combine) + chosen = ((selected[:, None] >> lanes[None, :].to(tl.uint64)) & 1) != 0 + tokens = words[:, None].to(tl.int32) * 64 + lanes[None, :] + logical_pages = tokens // TOKENS_PER_BLOCK + pages = tl.load( + page_base + tl.minimum(logical_pages, PAGE_TABLE_WIDTH - 1) * block_table_stride_page, + mask=chosen & (logical_pages < PAGE_TABLE_WIDTH), + other=-1, + ) # fmt: skip + kept = chosen & (pages >= 0) & (pages < NUM_CACHE_PAGES) + kept_count = kept.to(tl.int32) + count = tl.sum(kept_count, axis=1) + first = listed + tl.cumsum(count, axis=0) - count + rank = first[:, None] + tl.cumsum(kept_count, axis=1) - kept_count + tl.store( + slots + group * slots_stride_group + rank, + pages * TOKENS_PER_BLOCK + tokens % TOKENS_PER_BLOCK, + mask=kept, + ) # fmt: skip + resident = tl.reduce(tl.where(kept, lane_bits[None, :], 0), 1, _or_combine) + packed = _compress_bits(row_words, resident[None, :]) + shift = (first % 64).to(tl.uint64) + destination = packed_bits + members[:, None] * bits_stride_row + (first // 64)[None, :] + writes = live[:, None] & in_range[None, :] & (packed != 0) + tl.atomic_or(destination, (packed << shift[None, :]).to(tl.int64, bitcast=True), + mask=writes, sem="relaxed") # fmt: skip + # Bits past the destination word spill into the next one (never for shift 0). + spill = (packed >> 1) >> (63 - shift[None, :]) + tl.atomic_or(destination + 1, spill.to(tl.int64, bitcast=True), + mask=writes & ((first % 64) + count > 64)[None, :], sem="relaxed") # fmt: skip + listed += tl.sum(count, axis=0) + tiles = tl.cdiv(listed, BLOCK_N) + shared = one_request & (tiles * SHARED_TILE_COST <= own_tiles * 4) + tl.store(group_plan + group * 2, tl.where(shared, tiles, -1)) + tl.store(group_plan + group * 2 + 1, listed) + if not shared: + base = tl.atomic_add(lone_count, tl.sum(live.to(tl.int32), axis=0)) + tl.store(lone_rows + base + offsets, members.to(tl.int32), mask=live) + + +@triton.jit +def _rowtile_attention_kernel( + q, k_cache, v_cache, slots, packed_bits, group_plan, output, + softmax_scale, rows, + q_stride_row: tl.constexpr, q_stride_head: tl.constexpr, q_stride_dim: tl.constexpr, + k_stride_page: tl.constexpr, k_stride_head: tl.constexpr, + k_stride_token: tl.constexpr, k_stride_dim: tl.constexpr, + v_stride_page: tl.constexpr, v_stride_head: tl.constexpr, + v_stride_token: tl.constexpr, v_stride_dim: tl.constexpr, + slots_stride_group: tl.constexpr, bits_stride_row: tl.constexpr, + output_stride_row: tl.constexpr, output_stride_head: tl.constexpr, + output_stride_dim: tl.constexpr, + GROUP_SIZE: tl.constexpr, TOKENS_PER_BLOCK: tl.constexpr, ROWS_PER_TILE: tl.constexpr, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, + HEAD_DIM: tl.constexpr, RESCALE_THRESHOLD: tl.constexpr, NUM_STAGES: tl.constexpr, + WARP_SPECIALIZE: tl.constexpr, USE_PDL: tl.constexpr, +): # fmt: skip + """A group of rows of one request: walk the group's packed slot list in + 64-slot tiles, each tile's K/V loaded once for all rows and heads, each + row masked to its own packed bitmap word. Groups the plan did not take + exit here and are served row by row by `_rowtile_lone_kernel`.""" + group = tl.program_id(0).to(tl.int64) + kv_head = tl.program_id(1) + tiles = tl.load(group_plan + group * 2) + if tiles < 0: + if USE_PDL: + # The lone pass may still be running: exit only once it has, so a + # later launch on the stream sees every row written. + tl.extra.cuda.gdc_wait() + return + listed = tl.load(group_plan + group * 2 + 1) + slot_base = slots + group * slots_stride_group + m = tl.arange(0, BLOCK_M) + local_row = m // GROUP_SIZE + head = m % GROUP_SIZE + row = group * ROWS_PER_TILE + local_row + live = (local_row < ROWS_PER_TILE) & (row < rows) + safe_row = tl.where(live, row, 0) + + dim_offsets = tl.arange(0, HEAD_DIM) + query_heads = kv_head * GROUP_SIZE + head + q_tile = tl.load( + q + safe_row[:, None] * q_stride_row + query_heads[:, None] * q_stride_head + + dim_offsets[None, :] * q_stride_dim, + mask=live[:, None], + other=0.0, + ) # fmt: skip + accumulator = tl.zeros([BLOCK_M, HEAD_DIM], tl.float32) + # Finite, so a row with nothing yet attended rescales by exactly zero. + running_max = tl.full([BLOCK_M], -1.0e30, tl.float32) + running_sum = tl.zeros([BLOCK_M], tl.float32) + token_offsets = tl.arange(0, BLOCK_N) + # The packed bitmap as int32 halves: word 2t holds list entries [64t, 64t+32), 2t+1 the rest. + word_base = packed_bits + safe_row * bits_stride_row + in_high_half = token_offsets >= 32 + # Moves bit (n % 32) of a word to the sign bit, so a select is one compare. + to_sign = 31 - token_offsets % 32 + for tile in tl.range(0, tiles, num_stages=NUM_STAGES, warp_specialize=WARP_SPECIALIZE): + low = tl.load(word_base + 2 * tile, mask=live, other=0) + high = tl.load(word_base + 2 * tile + 1, mask=live, other=0) + entries = tile * BLOCK_N + token_offsets + usable = entries < listed + cache_slots = tl.load(slot_base + entries, mask=usable, other=0) + safe_pages = (cache_slots // TOKENS_PER_BLOCK).to(tl.int64) + token_in_page = cache_slots % TOKENS_PER_BLOCK + # K and V tiles share one [slots, dim] layout, so the slot list is read + # in one layout per tile; the QK product takes K transposed. + keys = tl.trans(tl.load( + k_cache + safe_pages[:, None] * k_stride_page + kv_head * k_stride_head + + token_in_page[:, None] * k_stride_token + dim_offsets[None, :] * k_stride_dim, + mask=usable[:, None], + other=0.0, + ).to(q_tile.dtype)) # fmt: skip + values = tl.load( + v_cache + safe_pages[:, None] * v_stride_page + kv_head * v_stride_head + + token_in_page[:, None] * v_stride_token + dim_offsets[None, :] * v_stride_dim, + mask=usable[:, None], + other=0.0, + ).to(q_tile.dtype) # fmt: skip + # A row's bits are set only on listed entries, so they already imply residency. + word = tl.where(in_high_half[None, :], high[:, None], low[:, None]) + attend = (word << to_sign[None, :]) < 0 + scores = tl.dot(q_tile, keys) * (softmax_scale * 1.4426950408889634) + scores = tl.where(attend, scores, -float("inf")) + next_max = tl.maximum(running_max, tl.max(scores, axis=1)) + if RESCALE_THRESHOLD > 0: + # Rescale only when some row's maximum grew past the threshold; + # otherwise exponentiate against the stale maximum (probabilities + # stay below 2**RESCALE_THRESHOLD; the final division is unchanged). + if tl.max(next_max - running_max, axis=0) > RESCALE_THRESHOLD: + correction = tl.math.exp2(running_max - next_max) + accumulator = accumulator * correction[:, None] + running_sum = running_sum * correction + running_max = next_max + else: + correction = tl.math.exp2(running_max - next_max) + accumulator = accumulator * correction[:, None] + running_sum = running_sum * correction + running_max = next_max + # Unattended scores are -inf, so their probabilities are exactly zero. + probabilities = tl.math.exp2(scores - running_max[:, None]) + accumulator = tl.dot(probabilities.to(q_tile.dtype), values, accumulator) + running_sum = running_sum + tl.sum(probabilities, axis=1) + + normalized = tl.where( + running_sum[:, None] > 0, accumulator / tl.maximum(running_sum[:, None], 1.0e-20), 0.0 + ) + tl.store( + output + safe_row[:, None] * output_stride_row + query_heads[:, None] * output_stride_head + + dim_offsets[None, :] * output_stride_dim, + normalized.to(output.dtype.element_ty), + mask=live[:, None], + ) # fmt: skip + if USE_PDL: + tl.extra.cuda.gdc_wait() + + +@triton.jit +def _rowtile_lone_kernel( + q, k_cache, v_cache, block_table, selected_tokens, request_indices, row_info, + lone_rows, lone_count, output, softmax_scale, rows, + q_stride_row: tl.constexpr, q_stride_head: tl.constexpr, q_stride_dim: tl.constexpr, + k_stride_page: tl.constexpr, k_stride_head: tl.constexpr, + k_stride_token: tl.constexpr, k_stride_dim: tl.constexpr, + v_stride_page: tl.constexpr, v_stride_head: tl.constexpr, + v_stride_token: tl.constexpr, v_stride_dim: tl.constexpr, + block_table_stride_request: tl.constexpr, block_table_stride_page: tl.constexpr, + selected_stride_row: tl.constexpr, selected_stride_token: tl.constexpr, + output_stride_row: tl.constexpr, output_stride_head: tl.constexpr, + output_stride_dim: tl.constexpr, + NUM_CACHE_PAGES: tl.constexpr, NUM_REQUESTS: tl.constexpr, PAGE_TABLE_WIDTH: tl.constexpr, + GROUP_SIZE: tl.constexpr, TOKENS_PER_BLOCK: tl.constexpr, ROWS_PER_TILE: tl.constexpr, + ROWS_POW2: tl.constexpr, BLOCK_H: tl.constexpr, BLOCK_N: tl.constexpr, + HEAD_DIM: tl.constexpr, LONE_PROGRAMS: tl.constexpr, USE_PDL: tl.constexpr, + LONE_STAGES: tl.constexpr, +): # fmt: skip + """Rows of groups the row-tile pass does not serve: each its own columns, gathered.""" + if USE_PDL: + # Independent of the row-tile pass (disjoint rows): let it start now. + tl.extra.cuda.gdc_launch_dependents() + kv_head = tl.program_id(1) + for slot in range(tl.program_id(0), tl.load(lone_count), LONE_PROGRAMS): + _rowtile_lone_row( + q, k_cache, v_cache, block_table, selected_tokens, request_indices, row_info, output, + softmax_scale, tl.load(lone_rows + slot).to(tl.int64), kv_head, + q_stride_row, q_stride_head, q_stride_dim, + k_stride_page, k_stride_head, k_stride_token, k_stride_dim, + v_stride_page, v_stride_head, v_stride_token, v_stride_dim, + block_table_stride_request, block_table_stride_page, + selected_stride_row, selected_stride_token, + output_stride_row, output_stride_head, output_stride_dim, + NUM_CACHE_PAGES, NUM_REQUESTS, PAGE_TABLE_WIDTH, GROUP_SIZE, TOKENS_PER_BLOCK, + BLOCK_H, BLOCK_N, HEAD_DIM, LONE_STAGES, + ) # fmt: skip + + +@triton.jit +def _rowtile_lone_row( + q, k_cache, v_cache, block_table, selected_tokens, request_indices, row_info, output, + softmax_scale, row, kv_head, + q_stride_row, q_stride_head, q_stride_dim, + k_stride_page, k_stride_head, k_stride_token, k_stride_dim, + v_stride_page, v_stride_head, v_stride_token, v_stride_dim, + block_table_stride_request, block_table_stride_page, + selected_stride_row, selected_stride_token, + output_stride_row, output_stride_head, output_stride_dim, + NUM_CACHE_PAGES: tl.constexpr, NUM_REQUESTS: tl.constexpr, PAGE_TABLE_WIDTH: tl.constexpr, + GROUP_SIZE: tl.constexpr, TOKENS_PER_BLOCK: tl.constexpr, BLOCK_H: tl.constexpr, + BLOCK_N: tl.constexpr, HEAD_DIM: tl.constexpr, LONE_STAGES: tl.constexpr, +): # fmt: skip + request = tl.load(request_indices + row).to(tl.int64) + valid_request = (request >= 0) & (request < NUM_REQUESTS) + safe_request = tl.minimum(tl.maximum(request, 0), NUM_REQUESTS - 1) + page_base = block_table + safe_request * block_table_stride_request + columns_read = tl.load(row_info + row * 2 + 1) + head_offsets = tl.arange(0, BLOCK_H) + dim_offsets = tl.arange(0, HEAD_DIM) + query_heads = kv_head * GROUP_SIZE + head_offsets + head_live = head_offsets < GROUP_SIZE + q_tile = tl.load( + q + row * q_stride_row + query_heads[:, None] * q_stride_head + + dim_offsets[None, :] * q_stride_dim, + mask=head_live[:, None], + other=0.0, + ) # fmt: skip + accumulator = tl.zeros([BLOCK_H, HEAD_DIM], tl.float32) + running_max = tl.full([BLOCK_H], -1.0e30, tl.float32) + running_sum = tl.zeros([BLOCK_H], tl.float32) + token_offsets = tl.arange(0, BLOCK_N) + for tile in tl.range(0, tl.cdiv(columns_read, BLOCK_N), num_stages=LONE_STAGES): + columns = tile * BLOCK_N + token_offsets + tokens = tl.load( + selected_tokens + row * selected_stride_row + columns * selected_stride_token, + mask=columns < columns_read, + other=-1, + ).to(tl.int32) + safe_tokens = tl.maximum(tokens, 0) + logical_pages = safe_tokens // TOKENS_PER_BLOCK + token_in_page = safe_tokens % TOKENS_PER_BLOCK + usable = valid_request & (tokens >= 0) & (logical_pages < PAGE_TABLE_WIDTH) + physical_pages = tl.load( + page_base + tl.minimum(logical_pages, PAGE_TABLE_WIDTH - 1) * block_table_stride_page, + mask=usable, + other=-1, + ).to(tl.int64) + usable = usable & (physical_pages >= 0) & (physical_pages < NUM_CACHE_PAGES) + safe_pages = tl.maximum(physical_pages, 0) + keys = tl.load( + k_cache + safe_pages[None, :] * k_stride_page + kv_head * k_stride_head + + token_in_page[None, :] * k_stride_token + dim_offsets[:, None] * k_stride_dim, + mask=usable[None, :], + other=0.0, + ).to(q_tile.dtype) # fmt: skip + values = tl.load( + v_cache + safe_pages[:, None] * v_stride_page + kv_head * v_stride_head + + token_in_page[:, None] * v_stride_token + dim_offsets[None, :] * v_stride_dim, + mask=usable[:, None], + other=0.0, + ).to(q_tile.dtype) # fmt: skip + scores = tl.dot(q_tile, keys) * (softmax_scale * 1.4426950408889634) + scores = tl.where(usable[None, :], scores, -float("inf")) + next_max = tl.maximum(running_max, tl.max(scores, axis=1)) + correction = tl.math.exp2(running_max - next_max) + probabilities = tl.where(usable[None, :], tl.math.exp2(scores - next_max[:, None]), 0.0) + accumulator = tl.dot( + probabilities.to(q_tile.dtype), values, accumulator * correction[:, None] + ) + running_sum = running_sum * correction + tl.sum(probabilities, axis=1) + running_max = next_max + normalized = tl.where( + running_sum[:, None] > 0, accumulator / tl.maximum(running_sum[:, None], 1.0e-20), 0.0 + ) + tl.store( + output + row * output_stride_row + query_heads[:, None] * output_stride_head + + dim_offsets[None, :] * output_stride_dim, + normalized.to(output.dtype.element_ty), + mask=head_live[:, None], + ) # fmt: skip + + +def _check_inputs(q, k_cache, v_cache, block_table, selected_tokens, request_indices, + tokens_per_block, softmax_scale, query_positions, compress_ratio): # fmt: skip + if q.ndim != 3: + raise ValueError("QSA queries must have shape [rows, query_heads, head_dim]") + if k_cache.ndim != 4 or v_cache.shape != k_cache.shape: + raise ValueError("QSA K/V caches must have matching HND paged shapes") + if block_table.ndim != 2 or block_table.dtype != torch.int32: + raise ValueError("QSA sparse GQA requires a two-dimensional int32 block table") + if selected_tokens.ndim != 2 or selected_tokens.dtype != torch.int32: + raise ValueError("QSA selected tokens must be a two-dimensional int32 tensor") + rows, num_q_heads, head_dim = q.shape + if selected_tokens.shape[0] != rows or selected_tokens.shape[1] <= 0: + raise ValueError("QSA selected-token rows must match non-empty query rows") + if request_indices.shape != (rows,) or request_indices.dtype not in _INTEGER_DTYPES: + raise ValueError("QSA request indices must have one integer entry per row") + if not request_indices.is_contiguous(): + raise ValueError("QSA request indices must be contiguous") + if tokens_per_block <= 0 or k_cache.shape[2] != tokens_per_block: + raise ValueError("QSA cache token dimension must match positive tokens_per_block") + if k_cache.shape[3] != head_dim or k_cache.dtype != v_cache.dtype: + raise ValueError("QSA query and K/V cache head geometry is incompatible") + if q.dtype not in (torch.bfloat16, torch.float16): + raise ValueError("QSA queries must use BF16 or FP16") + if k_cache.dtype not in (torch.bfloat16, torch.float16, torch.float8_e4m3fn): + raise ValueError("QSA K/V cache must use BF16, FP16, or FP8 E4M3") + if head_dim < 16 or head_dim & (head_dim - 1): + raise ValueError("QSA sparse GQA requires a power-of-two head dimension") + num_kv_heads = k_cache.shape[1] + if num_kv_heads <= 0 or num_q_heads <= 0 or num_q_heads % num_kv_heads: + raise ValueError("QSA query heads must be divisible by non-zero KV heads") + if k_cache.shape[0] <= 0 or block_table.shape[0] <= 0 or block_table.shape[1] <= 0: + raise ValueError("QSA sparse GQA requires non-empty cache and block-table geometry") + if k_cache.shape[0] * tokens_per_block >= 2**31: + raise ValueError("QSA row-tile cache slots must fit int32") + if not math.isfinite(softmax_scale) or softmax_scale <= 0: + raise ValueError("QSA sparse GQA softmax scale must be finite and positive") + if query_positions is None or compress_ratio is None: + raise ValueError("the row-tile kernel needs the prefill pair (positions, compress ratio)") + if query_positions.shape != (rows,) or query_positions.dtype not in _INTEGER_DTYPES: + raise ValueError("QSA query positions must be one integer entry per row") + if compress_ratio <= 0: + raise ValueError("QSA compression ratio must be positive") + token_topk = selected_tokens.shape[1] - compress_ratio + 1 + if token_topk <= 0 or token_topk % compress_ratio: + raise ValueError("QSA selected-token width is incompatible with its compression ratio") + devices = {t.device for t in (q, k_cache, v_cache, block_table, selected_tokens, + request_indices, query_positions)} # fmt: skip + if len(devices) != 1 or next(iter(devices)).type != "cuda": + raise ValueError("QSA row-tile sparse GQA needs every tensor on one CUDA device") + + +def qsa_rowtile_sparse_gqa( + *, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + block_table: torch.Tensor, + selected_tokens: torch.Tensor, + request_indices: torch.Tensor, + tokens_per_block: int, + softmax_scale: float, + query_positions: torch.Tensor, + compress_ratio: int, + block_m: int = 128, + block_n: int = 64, + num_warps: int = 8, + num_stages: int = 3, + rescale_threshold: float = 0.0, + warp_specialize: bool = False, + shared_tile_cost: int = 10, + lone_warps: int = 2, + lone_stages: int = 2, + plan_warps: int = 4, + plan_words: int = _PLAN_WORDS, +) -> torch.Tensor: + """Bitmap pass, plan pass, lone-row pass, row-tile pass; returns `[rows, Hq, D]` contiguous. + The row-tile pass runs `qsa_rowtile_attention_cute` where it is eligible (SM103, bf16 q and + K/V, D = 256, the default tile) and the Triton kernel otherwise; the Triton launch knobs + (`num_warps`, `num_stages`, `rescale_threshold`, `warp_specialize`) apply to the latter.""" + _check_inputs(q, k_cache, v_cache, block_table, selected_tokens, request_indices, + tokens_per_block, softmax_scale, query_positions, compress_ratio) # fmt: skip + rows, num_q_heads, head_dim = q.shape + num_kv_heads = k_cache.shape[1] + group_size = num_q_heads // num_kv_heads + if block_m < group_size or block_m & (block_m - 1): + raise ValueError("block_m must be a power of two holding at least one row's heads") + output = torch.empty((rows, num_q_heads, head_dim), dtype=q.dtype, device=q.device) + if rows == 0: + return output + table_width = block_table.shape[1] + if block_n != 64: + raise ValueError("block_n must be the 64-token bitmap word") + num_words = triton.cdiv(table_width * tokens_per_block, 64) + bits, row_info = _row_bits( + selected_tokens, query_positions, compress_ratio, tokens_per_block, table_width, num_words + ) + rows_per_tile = block_m // group_size + rows_pow2 = triton.next_power_of_2(rows_per_tile) + kv_strides = ( + k_cache.stride(0), k_cache.stride(1), k_cache.stride(2), k_cache.stride(3), + v_cache.stride(0), v_cache.stride(1), v_cache.stride(2), v_cache.stride(3), + ) # fmt: skip + q_strides = (q.stride(0), q.stride(1), q.stride(2)) + common = (*q_strides, *kv_strides, block_table.stride(0), block_table.stride(1)) + geometry = dict( + NUM_CACHE_PAGES=k_cache.shape[0], NUM_REQUESTS=block_table.shape[0], + PAGE_TABLE_WIDTH=table_width, GROUP_SIZE=group_size, + TOKENS_PER_BLOCK=tokens_per_block, ROWS_PER_TILE=rows_per_tile, ROWS_POW2=rows_pow2, + HEAD_DIM=head_dim, + ) # fmt: skip + out_strides = (output.stride(0), output.stride(1), output.stride(2)) + num_groups = triton.cdiv(rows, rows_per_tile) + slots = torch.empty((num_groups, num_words * 64), dtype=torch.int32, device=q.device) + packed_bits = torch.zeros((rows, num_words), dtype=torch.int64, device=q.device) + group_plan = torch.empty((num_groups, 2), dtype=torch.int32, device=q.device) + lone_rows = torch.empty((rows,), dtype=torch.int32, device=q.device) + lone_count = torch.zeros((1,), dtype=torch.int32, device=q.device) + _rowtile_plan_kernel[(num_groups,)]( + request_indices, row_info, bits, block_table, slots, packed_bits, group_plan, + lone_rows, lone_count, rows, + bits.stride(0), block_table.stride(0), block_table.stride(1), slots.stride(0), + ROWS_PER_TILE=rows_per_tile, NUM_REQUESTS=block_table.shape[0], ROWS_POW2=rows_pow2, + BLOCK_N=block_n, SHARED_TILE_COST=shared_tile_cost, NUM_CACHE_PAGES=k_cache.shape[0], + PAGE_TABLE_WIDTH=table_width, TOKENS_PER_BLOCK=tokens_per_block, + PLAN_WORDS=plan_words, num_warps=plan_warps, + ) # fmt: skip + use_pdl = _pdl_enabled() + _rowtile_lone_kernel[(min(rows, _LONE_PROGRAMS), num_kv_heads)]( + q, k_cache, v_cache, block_table, selected_tokens, request_indices, row_info, + lone_rows, lone_count, output, softmax_scale, rows, + *common, selected_tokens.stride(0), selected_tokens.stride(1), *out_strides, + BLOCK_H=max(16, triton.next_power_of_2(group_size)), BLOCK_N=64, + LONE_PROGRAMS=min(rows, _LONE_PROGRAMS), USE_PDL=use_pdl, LONE_STAGES=lone_stages, + num_warps=lone_warps, **geometry, + ) # fmt: skip + if _cute_eligible(q, k_cache, v_cache, block_m, block_n): + _rowtile_attention_cute[(num_groups, num_kv_heads)]( + q, k_cache, v_cache, slots, packed_bits, group_plan, output, softmax_scale, rows, + group_size=group_size, rows_per_tile=rows_per_tile, tokens_per_block=tokens_per_block, + use_pdl=use_pdl, + ) # fmt: skip + return output + _rowtile_attention_kernel[(num_groups, num_kv_heads)]( + q, k_cache, v_cache, slots, packed_bits.view(torch.int32), group_plan, output, + softmax_scale, rows, *q_strides, *kv_strides, slots.stride(0), 2 * packed_bits.stride(0), + *out_strides, + GROUP_SIZE=group_size, TOKENS_PER_BLOCK=tokens_per_block, ROWS_PER_TILE=rows_per_tile, + BLOCK_M=block_m, BLOCK_N=block_n, HEAD_DIM=head_dim, + RESCALE_THRESHOLD=rescale_threshold, NUM_STAGES=num_stages, + WARP_SPECIALIZE=warp_specialize, USE_PDL=use_pdl, + num_warps=num_warps, launch_pdl=use_pdl, + ) # fmt: skip + return output + + +_CHECKPOINT = dict( + q_heads=24, + kv_heads=2, + head_dim=256, + tokens_per_block=32, + compress_ratio=4, + token_topk=2048, + cache_dtype=torch.bfloat16, + q_dtype=torch.bfloat16, + # K and V are the two role planes of one [pages, 2, kv, tpb, D] pool. + cache_layout="pool_planes", + # q is the first 24 heads of the packed [T, 24 + 2 + 2, 256] QKV buffer. + q_layout="qkv_slice", + request_dtype=torch.int32, + # The engine's table at max_seq_len 10240: 320 pages per slot, -1 past each sequence. + pages=320, +) + + +def _spec(segments, slots=32, **overrides) -> dict: + """`segments` is `(first position, rows)` per request, in row order.""" + return {**_CHECKPOINT, "segments": tuple(segments), "slots": slots, **overrides} + + +def _decodes(count, rows_each=1, first=8200): + return [(first + (37 * g) % 1000, rows_each) for g in range(count)] + + +class QsaSparseGqaRowtile(OpWrapper): + """Sparse GQA over the selected tokens of an HND paged K/V pool, for a + launch that carries the prefill pair (`query_positions`, `compress_ratio`). + + For row `r`, query head `h`, KV head `h // (Hq // Hkv)`, over the columns + `j < bound(r)` of `selected_tokens[r]` that name a resident token:: + + t = selected_tokens[r, j] # logical position + page = block_table[request_indices[r], t // tokens_per_block] + s_j = (q[r, h] . k_cache[page, kv, t % tpb]) * softmax_scale + out[r, h] = sum_j softmax(s)_j * v_cache[page, kv, t % tpb] + + `bound(r)` is the per-row kernel's prefill bound -- `min(visible // cr, + block_topk) * cr + visible % cr` columns with `visible = query_positions[r] + + 1`, rounded up to 64 -- so a row reads exactly the columns + `qsa_paged_sparse_gqa` reads on a context launch. Negative indices, pages + past the table, unmapped or out-of-pool pages and foreign requests + contribute nothing; a row left with nothing returns zeros. Returns a + fresh, contiguous `[rows, Hq, D]` tensor of `q.dtype`. + + Four launches. A bitmap pass: one bit per selected logical token per row. + A plan pass per group of 10 consecutive rows of one request: the union of + the rows' selected tokens whose page is mapped, listed in logical order + as physical cache slots, and each row's bitmap re-indexed onto that list + (a per-word parallel bit extract). The group goes to the row-tile pass + when walking its list in 64-slot tiles costs at most 2.5x per-row tiles + less than the rows' own columns would; otherwise -- or when its rows span + requests -- its rows join the lone list. A lone-row pass (256 programs + per KV head striding the list) serves listed rows over their own columns, + as the per-row kernel does. The row-tile pass runs one program per + (group, KV head), loads each listed tile's K/V once for the 120 (row, + head) pairs and masks each row to its own re-indexed bits; it is launched + programmatically dependent on the lone pass, whose rows are disjoint, so + the two overlap. On SM103 with bf16 q and K/V and D = 256 that pass is + `qsa_rowtile_attention_cute` (tcgen05 MMAs, S / P / O in tensor memory, + a cp.async gather of the K/V rows), elsewhere a Triton kernel. + + Fusion boundary: as `qsa_paged_sparse_gqa` -- the block-table gather, + scale, online softmax and P.V are inside; deciding what to attend to, + appending K/V and q's projection, norm, RoPE and gate are not. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="context 300 rows from 0: rows see 1..300 tokens", spec=_spec([(0, 300)])), + Cell(why="context 272 rows from 2032: across the budget", spec=_spec([(2032, 272)])), + Cell( + why="two requests (150 @1900, 150 @900): a group spans them", + spec=_spec([(1900, 150), (900, 150)]), + ), + Cell(why="chunked continuation: 4096 rows from 4096", spec=_spec([(4096, 4096)])), + Cell( + why="sparse union: 4096 rows from 4096 ranking one score vector", + spec=_spec([(4096, 4096)], shared_scores=True), + ), + Cell(why="con 1 prefill: 8192 rows of one request from 0", spec=_spec([(0, 8192)])), + Cell(why="warm-up: 8191 rows of one request from 0", spec=_spec([(0, 8191)])), + Cell( + why="max_num_tokens: 2 requests x 8192 rows from 0", + spec=_spec([(0, 8192), (0, 8192)]), + ), + Cell( + why="con 128 mixed: 3 chunks (16240 rows) + 52 decodes", + spec=_spec([(0, 8192), (2192, 6000), (6144, 2048)] + _decodes(52), slots=128), + ), + Cell(why="mixed: 1 chunk (195) + 31 MTP decodes", spec=_spec([(0, 195)] + _decodes(31, 2))), + Cell( + why="mixed: 4 chunks (8134 rows) + 28 MTP decodes", + spec=_spec([(0, 8000), (907, 100), (1814, 30), (2721, 4)] + _decodes(28, 2)), + ), + Cell( + why="mixed: 9 chunks (3760 rows) + 23 decodes", + spec=_spec( + [ + (0, 1024), + (907, 900), + (1814, 700), + (2721, 500), + (3628, 300), + (4535, 200), + (5442, 100), + (6349, 30), + (7256, 6), + ] + + _decodes(23) + ), + ), # fmt: skip + Cell( + why="fp16 pool and fp16 q", + spec=_spec([(1900, 300)], cache_dtype=torch.float16, q_dtype=torch.float16), + ), + Cell( + why="fp8-e4m3 pool under a bf16 q, read unscaled", + spec=_spec([(1900, 300)], cache_dtype=torch.float8_e4m3fn), + ), + Cell( + why="int64 request indices and positions", + spec=_spec([(1000, 300)], request_dtype=torch.int64), + ), + Cell( + why="standalone contiguous K and V buffers, contiguous q", + spec=_spec([(900, 300)], cache_layout="contiguous", q_layout="contiguous"), + ), + Cell( + why="holes: all-padding row, foreign request, unmapped page", + spec=_spec([(900, 40), (950, 40)], holes=True), + ), + ) + + note = """ + Metadata consumed: none -- every input is an argument, the same arguments + `qsa_paged_sparse_gqa` takes, with the prefill pair required. What the + target passes: `qsa_block_table` (physical page ids), the expansion's + `-1`-padded logical selections, the layer pool's K/V role planes, and the + logical positions with the checkpoint's compress ratio. + + Preconditions (neither is checkable without reading the contents): + - Within the bound, a row names each token at most once. The row-tile + pass attends to the *set* of tokens a row names; the per-row kernel + attends to every column, so a repeated index would weigh twice there + and once here. QSA's top-k + expansion never repeats one (distinct + blocks, and a tail outside every complete block): measured zero + repeats on four layers of the target's own 8192-token prefill. + Rows served by the lone-row pass keep per-column semantics. + - As for the per-row kernel: each row is prefix-packed with exactly the + bound's count of valid entries; entries past the 64-rounded bound + are ignored. + + Work: a row-tile group walks its packed list, so its MMA work grows + with the union of its rows' selections rather than with the group's + highest position. On the target's own selections (8192-token prefill, + layers 0/3/6/11) that union is 13-32% shorter than every 64-token tile up + to the highest token, and the walk loads ~5x fewer K/V tiles than the + per-row kernel. Measured on GB300 (eager, whole call) at 8192 context + rows of real selections: 1544-1565 us on layers 0-6 and 1274 us on layer + 11, against 1711-1715 us for a walk of every tile up to the highest token + and 2953 us for `qsa_paged_sparse_gqa`. Uniformly random selections + leave almost no gap in the union, and the call costs what the full walk + does (16292-row mixed step: 3859 vs 3836 us). Below a few hundred rows + it is slower than the per-row kernel (generation-only launches 1.4-2x); + the campaign's context launches are thousands of rows. The tensor-core + row-tile pass alone, on two dumped 16384-row context launches of the + target (GB300, same plan): 631 and 887 us against 1597 and 2313 us for + the Triton pass. + + Scratch per call: the bitmap and the re-indexed bitmap, each `rows x + ceil(table_width * tpb / 64)` int64 words (21 MB at 16384 rows and 320 + pages), the slot lists, `ceil(rows / 10) x` that many words' tokens as + int32 (67 MB at 16384 rows), the plan and the lone list. Calls allocate, + so capture in a CUDA graph is not certified; this entry serves the + prefill path, which runs eager. Results are deterministic: the bitmap + OR, the re-indexed bitmap OR and the lone list order do not change any + row's arithmetic. + """ + + def __call__( + self, + *, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + block_table: torch.Tensor, + selected_tokens: torch.Tensor, + request_indices: torch.Tensor, + tokens_per_block: int, + softmax_scale: float, + query_positions: torch.Tensor, + compress_ratio: int, + ) -> torch.Tensor: + return qsa_rowtile_sparse_gqa( + q=q, + k_cache=k_cache, + v_cache=v_cache, + block_table=block_table, + selected_tokens=selected_tokens, + request_indices=request_indices, + tokens_per_block=tokens_per_block, + softmax_scale=softmax_scale, + query_positions=query_positions, + compress_ratio=compress_ratio, + ) + + def reference( + self, + *, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + block_table: torch.Tensor, + selected_tokens: torch.Tensor, + request_indices: torch.Tensor, + tokens_per_block: int, + softmax_scale: float, + query_positions: torch.Tensor, + compress_ratio: int, + ) -> torch.Tensor: + """The reference `qsa_paged_sparse_gqa` certifies against: fp32 sparse + GQA over every resident selected token, rounded once. The prefill pair + is ignored, as there: on a prefix-packed selection the bound reads + exactly the valid entries.""" + return _per_row_entry.reference( + q=q, + k_cache=k_cache, + v_cache=v_cache, + block_table=block_table, + selected_tokens=selected_tokens, + request_indices=request_indices, + tokens_per_block=tokens_per_block, + softmax_scale=softmax_scale, + query_positions=query_positions, + compress_ratio=compress_ratio, + ) + + +qsa_sparse_gqa_rowtile = QsaSparseGqaRowtile() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_expand_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_expand_decode.py new file mode 100644 index 000000000000..42be668f474f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_expand_decode.py @@ -0,0 +1,590 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""QSA's block top-k and its expansion to token indices, one launch per batch of rows. + +Tier 5 (CuTe DSL written for this target). The decode path ran `qsa_topk_select` +(the CUDA indexer top-k, one CTA a row) and then `qsa_expand_block_indices` (one +more launch) per attention layer. Here one CTA per row holds four consecutive scores a +thread and narrows to the k-th largest by radix passes over the scores' bit patterns +(shared-memory histograms of 11 bits), starting below the bits every visible score +shares. It stops as soon as the bin holding the threshold has at most `CAP` keys and +ranks those exactly by pairwise comparison (ties to the higher block, the canonical +selection's rule). Every selected block is then placed in block order by one block-wide +scan, and expanded to token indices with the causal tail, as the expansion entry does. + +Scores are a sum of ReLUs times a positive scale, so non-negative: their fp32 bit +patterns, read as int32, order like their values (a negative zero is taken as zero, a +NaN ranks above everything). +""" + +from __future__ import annotations + +import functools +import os + +import cuda.bindings.driver as cuda_drv +import cutlass +import cutlass.cute as cute +import torch +from cutlass import const_expr +from cutlass.cute import experimental as cute_ext +from cutlass.cute.runtime import from_dlpack + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .qsa_expand_block_indices import qsa_expand_block_indices as _expand + +#: Consecutive scores a thread holds (one 16-byte load). +KPT = 4 +#: CTAs an SM the register allocation is held to. At the 2592-wide score table a row is 672 +#: threads, which the per-subpartition register file holds once an SM at the compiler's own +#: 48 registers a thread; at three (32 a thread) rows overlap each other's barriers and +#: shared-memory round trips (GB300, a mixed step's 16256 context rows: 426 -> 236 us on +#: half-zero rows, 486 -> 284 on relu-sparse ones; 128 decode rows 8.2 us both ways). +MIN_BLOCKS = 3 +#: Radix passes at most (11 bits each; 31 magnitude bits). +_PASSES = 3 +_HIST = 2048 +_I32_MAX = 0x7FFFFFFF +#: At most this many keys in the threshold's bin are ranked pairwise instead of +#: narrowed by another pass. +CAP = 128 +#: Threads that own the 2048 bins of a scan, four each. +_SCAN_THREADS = _HIST // 4 + + +@functools.lru_cache(maxsize=None) +def pdl_enabled(device_index: int) -> bool: + """Programmatic dependent launch: sm_90+ unless TRTLLM_ENABLE_PDL=0.""" + props = torch.cuda.get_device_properties(device_index) + return props.major >= 9 and os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" + + +def _threads(width: int) -> int: + return max(_SCAN_THREADS, -(-width // (KPT * 32)) * 32) + + +class _TopkExpandKernel: + def __init__(self, *, width, block_topk, compress, token_topk, pdl): + self.width, self.k, self.cr, self.token_topk, self.pdl = ( + width, + block_topk, + compress, + token_topk, + pdl, + ) + self.final = token_topk + compress - 1 + self.threads = _threads(width) + self.warps = self.threads // 32 + self.flag_words = -(-width // 32) + + @cute.jit + def __call__(self, mKeys, mEnd, mPos, mLen, mTopk, mSel, stream: cuda_drv.CUstream): + rows = cute.size(mKeys, mode=[0]) + ld_keys = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), cutlass.Int32, num_bits_per_copy=32 * KPT + ) + self.kernel(mKeys, mEnd, mPos, mLen, mTopk, mSel, ld_keys).launch( + grid=[rows, 1, 1], + block=[self.threads, 1, 1], + smem=(_HIST + self.flag_words + 2 * CAP + 64 + 16) * 4 + 128, + stream=stream, + use_pdl=self.pdl, + min_blocks_per_mp=MIN_BLOCKS, + ) + + @cute.kernel + def kernel(self, mKeys, mEnd, mPos, mLen, mTopk, mSel, ld_keys: cute.CopyAtom): + tidx, _, _ = cute.arch.thread_idx() + row, _, _ = cute.arch.block_idx() + warp = cute.arch.warp_idx() + lane = cute.arch.lane_idx() + nt: cutlass.Constexpr = self.threads + nw: cutlass.Constexpr = self.warps + k: cutlass.Constexpr = self.k + cr: cutlass.Constexpr = self.cr + smem = cutlass.utils.SmemAllocator() + hist = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((_HIST,), stride=(1,)), byte_alignment=16 + ) + flags = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((self.flag_words,), stride=(1,)), byte_alignment=16 + ) + cand_key = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((CAP,), stride=(1,)), byte_alignment=16 + ) + cand_idx = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((CAP,), stride=(1,)), byte_alignment=16 + ) + wa = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((32,), stride=(1,)), byte_alignment=16 + ) + wb = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((32,), stride=(1,)), byte_alignment=16 + ) + state = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((16,), stride=(1,)), byte_alignment=16 + ) + + # 0. Shared state, before the wait (nothing global is read here). + for b in cutlass.range_constexpr(-(-_HIST // nt)): + if tidx + b * nt < _HIST: + hist[tidx + b * nt] = cutlass.Int32(0) + for b in cutlass.range_constexpr(-(-self.flag_words // nt)): + if tidx + b * nt < self.flag_words: + flags[tidx + b * nt] = cutlass.Int32(0) + if tidx == 0: + state[5] = cutlass.Int32(0) # candidates compacted + if const_expr(self.pdl): + cute.arch.griddepcontrol_wait() + cute.arch.griddepcontrol_launch_dependents() + + # 1. The row's metadata and four consecutive keys a thread (a column past the + # window is no candidate: -1); the largest and smallest key in the window. + n = mEnd[row] + seqlen = mLen[row] + visible = cutlass.Int32(mPos[row]) + 1 + e0 = tidx * KPT + keys = cute.make_rmem_tensor(cute.make_layout((KPT,), stride=(1,)), cutlass.Int32) + keys.fill(-1) + if e0 < self.width: + cute.copy(ld_keys, cute.logical_divide(mKeys, (None, KPT))[row, (None, tidx)], keys) + kmax = cutlass.Int32(-1) + kmin = cutlass.Int32(_I32_MAX) + for c in cutlass.range_constexpr(KPT): + b = keys[c] + if b < 0: # -0.0 (or a negative): rank as zero + b = cutlass.Int32(0) + if e0 + c >= n: + b = cutlass.Int32(-1) + keys[c] = b + kmax = cutlass.max(kmax, b) + if b >= 0: + kmin = cutlass.min(kmin, b) + kmax = cute.arch.warp_reduction_max(kmax) + kmin = cute.arch.warp_reduction(kmin, lambda x, y: cutlass.min(x, y)) + if lane == 0: + wa[warp] = kmax + wb[warp] = kmin + cute.arch.sync_threads() + gmax = cutlass.Int32(-1) + gmin = cutlass.Int32(_I32_MAX) + if lane < nw: + gmax = wa[lane] + gmin = wb[lane] + gmax = cute.arch.warp_reduction_max(gmax) + gmin = cute.arch.warp_reduction(gmin, lambda x, y: cutlass.min(x, y)) + + # 2. Every key in the window shares the bits above the highest bit where the + # largest and the smallest differ (`hb`); radix passes walk down from there, 11 bits + # at a time, until the threshold's bin holds at most CAP keys or no bits are left. + diff = gmax ^ gmin + hb = cutlass.Int32(-1) + if diff != 0: + hb = 31 - cutlass.Int32(cute.arch.clz(diff)) + low = (cutlass.Int64(1) << cutlass.Int64(hb + 1)) - 1 + mask0 = cutlass.Int32(cutlass.Int64(_I32_MAX) & (cutlass.Int64(-1) ^ low)) + if tidx == 0: + state[0] = gmax & mask0 # threshold prefix + state[1] = mask0 # its mask + state[2] = cutlass.Int32(k) # still needed at the threshold + state[3] = n # keys in the threshold's bin + state[4] = cutlass.Int32(1) # narrowing done + state[6] = cutlass.max(hb - 10, 0) # the next pass's shift + if n > k: + if hb >= 0: + state[4] = cutlass.Int32(0) + cute.arch.sync_threads() + for _p in cutlass.range_constexpr(_PASSES): + if state[4] == 0: + prefix = state[0] + mask = state[1] + need = state[2] + shift = state[6] + for c in cutlass.range_constexpr(KPT): + if keys[c] >= 0: + if (keys[c] & mask) == prefix: + digit = (keys[c] >> shift) & (_HIST - 1) + cute.arch.atomic_add( + hist.iterator + digit, cutlass.Int32(1), scope="cta" + ) + cute.arch.sync_threads() + # Thread t < 512 owns bins 2047-4t .. 2044-4t, counted from the top; each is + # cleared for the next pass once read. + h = cute.make_rmem_tensor(cute.make_layout((4,), stride=(1,)), cutlass.Int32) + h.fill(0) + own = cutlass.Int32(0) + if tidx < _SCAN_THREADS: + for j in cutlass.range_constexpr(4): + h[j] = hist[_HIST - 1 - 4 * tidx - j] + hist[_HIST - 1 - 4 * tidx - j] = cutlass.Int32(0) + own = own + h[j] + incl = own + for off in (1, 2, 4, 8, 16): + up = cute.arch.shuffle_sync_up(incl, offset=off) + if lane >= off: + incl = incl + up + if lane == 31: + wa[warp] = incl + cute.arch.sync_threads() + above = cutlass.Int32(0) + if lane < warp: + above = wa[lane] + before = cute.arch.warp_reduction_sum(above) + incl - own # keys in the bins above + if tidx < _SCAN_THREADS: + if before < need: + if need <= before + own: + run = before + for j in cutlass.range_constexpr(4): + if run < need: + if need <= run + h[j]: + state[0] = prefix | ((_HIST - 1 - 4 * tidx - j) << shift) + state[1] = mask | ((_HIST - 1) << shift) + state[2] = need - run + state[3] = h[j] + state[4] = cutlass.Int32( + cutlass.select_((h[j] <= CAP) | (shift == 0), 1, 0) + ) + state[6] = cutlass.max(shift - 11, 0) + run = run + h[j] + cute.arch.sync_threads() + + prefix = state[0] + mask = state[1] + need = state[2] + at_thr = state[3] + # 3. At most CAP keys in the threshold's bin: rank them exactly, ties to the + # higher block, and flag the `need` best. + if n > k: + if at_thr <= CAP: + for c in cutlass.range_constexpr(KPT): + if keys[c] >= 0: + if (keys[c] & mask) == prefix: + slot = cute.arch.atomic_add( + state.iterator + 5, cutlass.Int32(1), scope="cta" + ) + cand_key[slot] = keys[c] + cand_idx[slot] = e0 + c + cute.arch.sync_threads() + if tidx < at_thr: + kt = cand_key[tidx] + it = cand_idx[tidx] + rank = cutlass.Int32(0) + for j in cutlass.range(0, at_thr, 1, unroll=8): + kj = cand_key[j] + ij = cand_idx[j] + better = (kj > kt) | ((kj == kt) & (ij > it)) + rank = rank + cutlass.Int32(cutlass.select_(better, 1, 0)) + if rank < need: + cute.arch.atomic_or( + flags.iterator + (it >> 5), cutlass.Int32(1) << (it & 31), scope="cta" + ) + cute.arch.sync_threads() + + # 4. Block order: one scan of (selected, tied at the threshold) counts. A tied key + # is left only when the passes ran out of bits with more than CAP keys tied; the + # last `need` of those in block order are taken. + take = cute.make_rmem_tensor(cute.make_layout((KPT,), stride=(1,)), cutlass.Int32) + tied = cute.make_rmem_tensor(cute.make_layout((KPT,), stride=(1,)), cutlass.Int32) + own = cutlass.Int32(0) + for c in cutlass.range_constexpr(KPT): + t_c = cutlass.Int32(0) + q_c = cutlass.Int32(0) + if keys[c] >= 0: + if n <= k: + t_c = cutlass.Int32(1) + else: + mk = keys[c] & mask + if mk > prefix: + t_c = cutlass.Int32(1) + if mk == prefix: + if at_thr > CAP: + q_c = cutlass.Int32(1) + else: + e = e0 + c + t_c = (flags[e >> 5] >> (e & 31)) & 1 + take[c] = t_c + tied[c] = q_c + own = own + t_c + (q_c << 16) + incl = own + for off in (1, 2, 4, 8, 16): + up = cute.arch.shuffle_sync_up(incl, offset=off) + if lane >= off: + incl = incl + up + if lane == 31: + wb[warp] = incl + cute.arch.sync_threads() + above = cutlass.Int32(0) + if lane < warp: + above = wb[lane] + pre = cute.arch.warp_reduction_sum(above) + incl - own + taken_before = pre & 0xFFFF + tied_before = pre >> 16 + # The tied keys left out, the lowest-index ones (none unless the threshold's bin + # holds more keys than are still needed). + skip = cutlass.max(at_thr - need, 0) + for c in cutlass.range_constexpr(KPT): + e = e0 + c + sel = cutlass.Int32(0) + pos = cutlass.Int32(0) + if tied[c] == 1: + if tied_before >= skip: + sel = cutlass.Int32(1) + pos = taken_before + tied_before - skip + tied_before = tied_before + 1 + if take[c] == 1: + sel = cutlass.Int32(1) + pos = taken_before + cutlass.max(tied_before - skip, 0) + taken_before = taken_before + 1 + if sel == 1: + mTopk[row, pos] = e + for j in cutlass.range_constexpr(cr): + tok = e * cr + j + val = cutlass.Int32(-1) + if tok < seqlen: + val = tok + mSel[row, pos * cr + j] = val + + # 5. Padding of the block list; the causal tail after the selected tokens, then -1. + total = n + if n > k: + total = cutlass.Int32(k) + for c in cutlass.range_constexpr(-(-k // nt)): + s = tidx + c * nt + if s < k: + if s >= total: + mTopk[row, s] = cutlass.Int32(-1) + tail_start = (visible // cr) * cr + tail_count = visible - tail_start + body = total * cr + if body > self.token_topk: + body = cutlass.Int32(self.token_topk) + for c in cutlass.range_constexpr(-(-self.final // nt)): + col = tidx + c * nt + if col < self.final: + if col >= body: + off = col - body + val = cutlass.Int32(-1) + if off < cr - 1: + if off < tail_count: + if tail_start + off < seqlen: + val = tail_start + off + mSel[row, col] = val + + +def _dl(t: torch.Tensor, align: int = 4): + return from_dlpack(t, assumed_align=align) + + +@functools.lru_cache(maxsize=None) +def _compiled(rows, width, key_stride, block_topk, compress, token_topk, pos_dtype, len_dtype, pdl): + kern = _TopkExpandKernel( + width=width, block_topk=block_topk, compress=compress, token_topk=token_topk, pdl=pdl + ) + dev = "cuda" + keys = torch.empty(rows, key_stride, dtype=torch.int32, device=dev)[:, :width] + end = torch.empty(rows, dtype=torch.int32, device=dev) + pos = torch.empty(rows, dtype=pos_dtype, device=dev) + ln = torch.empty(rows, dtype=len_dtype, device=dev) + topk = torch.empty(rows, block_topk, dtype=torch.int32, device=dev) + sel = torch.empty(rows, token_topk + compress - 1, dtype=torch.int32, device=dev) + stream = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream) + return cute_ext.compile( + kern, _dl(keys, 16), _dl(end), _dl(pos), _dl(ln), _dl(topk), _dl(sel), stream + ) + + +def launch( + scores, + row_ends, + query_positions, + sequence_lengths, + topk_indices, + compress_ratio, + token_topk, + pdl=True, + out=None, +): + """The launch behind `QsaTopkExpandDecode.__call__`: writes `topk_indices` and returns the + token selection `[rows, token_topk + compress_ratio - 1]` int32 (`out` when given).""" + rows, width = scores.shape + block_topk = topk_indices.shape[1] + dev = scores.device + sel = out + if sel is None: + sel = torch.empty((rows, token_topk + compress_ratio - 1), dtype=torch.int32, device=dev) + if rows == 0: + return sel + pdl = pdl and pdl_enabled(dev.index if dev.index is not None else torch.cuda.current_device()) + fn = _compiled( + rows, + width, + scores.stride(0), + block_topk, + compress_ratio, + token_topk, + query_positions.dtype, + sequence_lengths.dtype, + bool(pdl), + ) + stream = cuda_drv.CUstream(torch.cuda.current_stream(dev).cuda_stream) + fn( + _dl(scores.view(torch.int32), 16), + _dl(row_ends), + _dl(query_positions), + _dl(sequence_lengths), + _dl(topk_indices), + _dl(sel), + stream, + ) + return sel + + +class QsaTopkExpandDecode(OpWrapper): + """`qsa_topk_select` (zero row starts, `row_ends` the visible block counts), + `qsa_canonical_selection` and `qsa_expand_block_indices`, in one launch: one CTA a row, for + decode rows and a context chunk's rows alike. + + Per row: the `block_topk` largest of `scores[row, :row_ends[row]]` (non-negative fp32), + ties at the threshold to the higher block, written to `topk_indices[row]` in block order + and padded with -1 -- the composition's selection after `qsa_canonical_selection`; the + returned token selection is the expansion entry's: each selected group's members below + `sequence_lengths[row]`, packed at the front, then the incomplete group at the causal + frontier (`query_positions[row] + 1` visible tokens), then -1. `topk_indices` is written + in place, and the selection into `out` when given (a contiguous int32 + `[rows, token_topk + compress_ratio - 1]`, e.g. a row range of the step's whole + selection); nothing else is written. + + Fusion boundary: the selection and the expansion. Scoring is the caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode, 1 row at position 8700: 2175 visible blocks of 2560", + spec=dict(rows=1, pos=(8700,)), + ), + Cell( + why="decode, 4 rows across the budget, one below it", + spec=dict(rows=4, pos=(8700, 4095, 1500, 9900)), + ), + Cell( + why="decode, 1 row with ties at the threshold", + spec=dict(rows=1, pos=(8700,), ties=True), + ), + Cell( + why="decode, 2 rows, 2 MTP-style positions sharing a request", + spec=dict(rows=2, pos=(6001, 6002)), + ), + Cell( + why="decode, 1 row of scores within 5% of each other", + spec=dict(rows=1, pos=(8700,), narrow=True), + ), + Cell( + why="decode, 1 row of equal scores: block order decides", + spec=dict(rows=1, pos=(8700,), const=True), + ), + Cell( + why="decode, 1 row, 150 scores tied at the threshold", + spec=dict(rows=1, pos=(8700,), tied=150), + ), + ) + + note = """ + Metadata consumed: `qsa_visible_blocks` (`row_ends`), the rows' logical positions + and sequence lengths. Preconditions: zero row starts; `scores` fp32 `[rows, W]`, unit + column stride, 16-byte aligned rows, `W % 4 == 0` and `W <= 4096` (four keys a thread, + at most 1024 threads); `topk_indices` int32 `[rows, K]` contiguous. Ties: the higher + block at the threshold, as `qsa_canonical_selection` rewrites the CUDA top-k's pick. + """ + + def __call__( + self, + scores, + row_ends, + query_positions, + sequence_lengths, + topk_indices, + compress_ratio, + token_topk, + out=None, + ): + return launch( + scores, + row_ends, + query_positions, + sequence_lengths, + topk_indices, + compress_ratio, + token_topk, + out=out, + ) + + def reference( + self, + scores, + row_ends, + query_positions, + sequence_lengths, + topk_indices, + compress_ratio, + token_topk, + out=None, + ): + """`(topk_indices_after, selection)`: a stable descending sort per row (ties to the + higher block), the k first in block order, then the expansion entry's reference.""" + rows, _ = scores.shape + k = topk_indices.shape[1] + out = torch.full_like(topk_indices, -1) + for r in range(rows): + n = int(row_ends[r]) + s = scores[r, :n].float().clamp(min=0.0) + # A stable sort of the reversed row puts the higher block first among equals. + order = (n - 1) - torch.sort(torch.flip(s, [0]), descending=True, stable=True).indices[ + :k + ] + chosen = torch.sort(order).values + out[r, : chosen.numel()] = chosen.to(torch.int32) + sel = _expand.reference( + out, + query_positions, + sequence_lengths, + compress_ratio=compress_ratio, + token_topk=token_topk, + ) + return out, sel + + def is_valid( + self, + scores, + row_ends, + query_positions, + sequence_lengths, + topk_indices, + compress_ratio, + token_topk, + out=None, + ): + assert scores.dtype == torch.float32 and scores.stride(1) == 1, ( + "fp32 scores, unit column stride" + ) + assert scores.shape[1] % KPT == 0 and scores.shape[1] <= KPT * 1024, "W % 4 == 0, W <= 4096" + assert scores.stride(0) % KPT == 0 and scores.data_ptr() % 16 == 0, "16-byte aligned rows" + assert topk_indices.dtype == torch.int32 and topk_indices.is_contiguous(), ( + "int32 topk_indices" + ) + assert row_ends.dtype == torch.int32 and row_ends.shape == (scores.shape[0],), ( + "int32 row_ends [rows]" + ) + assert token_topk == topk_indices.shape[1] * compress_ratio, ( + "token_topk = k * compress_ratio" + ) + assert out is None or ( + out.dtype == torch.int32 + and out.is_contiguous() + and out.shape == (scores.shape[0], token_topk + compress_ratio - 1) + ), "out: contiguous int32 [rows, token_topk + compress_ratio - 1]" + + +qsa_topk_expand_decode = QsaTopkExpandDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_select.py index fda108df6ac2..a520d94dc9b9 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_select.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_select.py @@ -17,6 +17,11 @@ #: Score-row width the target hands over: a 260-page block table at 32 tokens #: per block, compressed 4:1. _SCORE_COLUMNS = 260 * 32 // 4 +#: Rows from which a contiguous launch runs the CuTe DSL sibling: the CUDA op +#: serves rows at or past 12288 with its radix variant, which is the slower of +#: the two (16384 rows: 310 vs 135 us at k=512 over 2592 columns, GB300); +#: below it the two run alike and the CUDA op needs no JIT compilation. +_CUTE_DSL_MIN_ROWS = 12288 class TopkExpectation(NamedTuple): @@ -41,7 +46,9 @@ class QsaTopkSelect(OpWrapper): """Each row's top-`index_topk` window-local column offsets, written into `indices`. Wraps `torch.ops.trtllm.indexer_topk_prefill` (CUDA, compiled into the - wheel). For row `r` it selects the `index_topk` largest values of + wheel) and, for a contiguous launch of `_CUTE_DSL_MIN_ROWS` rows or more, + its CuTe DSL sibling `torch.ops.trtllm.cute_dsl_indexer_topk_prefill_blackwell` + (the same contract). For row `r` it selects the `index_topk` largest values of `logits[r, row_starts[r]:row_ends[r]]` and writes their offsets *relative to `row_starts[r]`* into `indices[r]`: `min(index_topk, window)` distinct offsets as a prefix, `-1` in every remaining slot. Every slot of every row @@ -204,7 +211,18 @@ class QsaTopkSelect(OpWrapper): ), ), Cell( - why="row tier: 12416 rows, rows >= 12288 take the radix variant", + why="context: 16384 rows at positions 0..16383, NaN tail (CuTe DSL op)", + spec=dict( + rows=16384, + cols=_SCORE_COLUMNS, + index_topk=_BLOCK_TOPK, + windows=("chunk", 0), + logits="relu", + tail="nan", + ), + ), + Cell( + why="row tier: 12416 rows, past 12287 the CuTe DSL op", spec=dict( rows=6 * 2048 + 128, cols=512, @@ -235,16 +253,15 @@ class QsaTopkSelect(OpWrapper): 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 + The CUDA op dispatches by row index, not column count: rows below 12288 run + an insertion-sort variant, rows at or above a radix variant, in one launch. + A launch of `_CUTE_DSL_MIN_ROWS` (12288) rows or more with contiguous + logits runs the CuTe DSL sibling instead -- QSA's stock indexer's choice on + this GPU, measured to select identical sets (never identical permutations; + ties, as in the CUDA op, broken arbitrarily). The tuned config's + `max_num_tokens` 16384 reaches it at prefill-heavy mixed steps; smaller + launches, every decode one among them, keep the compiled CUDA op, so a + first captured decode JIT-compiles nothing. `torch.ops.trtllm.indexer_topk_decode` is the request-grouped sibling; QSA does not use it. CUDA graphs: captured at `[8, 2048]` and `[256, 2048]` (32 requests at @@ -253,7 +270,7 @@ class QsaTopkSelect(OpWrapper): replays is not a valid check (order is unspecified). Hard limits: not probed. `indices` is addressed as `row * index_topk + - slot`; the largest certified extent is 12416 x 64. + slot`; the largest certified extent is 16384 x 512. """ def __call__( @@ -264,7 +281,12 @@ def __call__( indices: torch.Tensor, index_topk: int = 2048, ) -> torch.Tensor: - torch.ops.trtllm.indexer_topk_prefill(logits, row_starts, row_ends, indices, index_topk) + if logits.shape[0] >= _CUTE_DSL_MIN_ROWS and logits.is_contiguous(): + torch.ops.trtllm.cute_dsl_indexer_topk_prefill_blackwell( + logits, row_starts, row_ends, indices, index_topk + ) + else: + torch.ops.trtllm.indexer_topk_prefill(logits, row_starts, row_ends, indices, index_topk) return indices def reference( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_attn_decode_prologue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_attn_decode_prologue.py new file mode 100644 index 000000000000..97f59f172d53 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_attn_decode_prologue.py @@ -0,0 +1,576 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A sparse decode attention layer's work between its input projection and its index +scores, one PDL launch. + +The composition it replaces ran four launches over two streams: the QK norm + rope of +both planes (`qwen4_exp_qk_norm_rope`), the paged K/V append (`qsa_paged_kv_store`), and on +the side stream the index plane's raw-key and coordinate store (`qsa_paged_index_store`) +and the group compression (`qsa_prefill_compress`), then a cross-stream join. Here the +programs of `qwen4_exp_qk_norm_rope` run as they do there; the ones holding k heads also +write their rotated rows into the K pool; one more program per row stores v, the index +key and the coordinates, and -- at a group's last token -- pools the group with this +token's key taken from the projection row (the cached copy is the one being written), +normalizes, rotates and stores it at the anchor slot, as the compression kernel does. + +With PDL every operand that is static or host-filled -- the page and slot, the rotation +factors at the row's positions, the norm weights, and at a group end the group's earlier +keys from the index plane with their rotation -- is read before `griddepcontrol.wait`; +after it a program reads only the projection row. +""" + +from __future__ import annotations + +import os +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .qsa_paged_index_store import qsa_paged_index_store as _index_store +from .qsa_paged_kv_store import qsa_paged_kv_store as _kv_store +from .qsa_prefill_compress import qsa_prefill_compress as _compress +from .qwen4_exp_qk_norm_rope import _launch as _qk_launch +from .qwen4_exp_qk_norm_rope import qwen4_exp_qk_norm_rope as _qk_norm_rope + + +@triton.jit +def _rope_factors(positions, num_tokens, token, cos_sin, stride_cs, D: tl.constexpr, ROT: tl.constexpr, + S1: tl.constexpr, S2: tl.constexpr): # fmt: skip + """A D-wide head's rotation factors at the row's positions (interleaved mRoPE: pair j + reads coordinate 1 when j % 3 == 1 and j < 3 * S1, coordinate 2 when j % 3 == 2 and + j < 3 * S2, else 0): cos, and sin signed for the NeoX halves; identity past ROT.""" + HALF: tl.constexpr = ROT // 2 + d = tl.arange(0, D) + rot = d < ROT + j = d % HALF + p0 = tl.load(positions + token) + p1 = tl.load(positions + num_tokens + token) + p2 = tl.load(positions + 2 * num_tokens + token) + pos = tl.where((j % 3 == 1) & (j < 3 * S1), p1, tl.where((j % 3 == 2) & (j < 3 * S2), p2, p0)) + row = cos_sin + pos.to(tl.int64) * stride_cs + c = tl.load(row + j, mask=rot, other=1.0) + s = tl.load(row + HALF + j, mask=rot, other=0.0) + s = tl.where(d < HALF, -s, s) + return c, s + + +@triton.jit +def _norm_weights( + heads, weight, weight_k, D: tl.constexpr, ROT: tl.constexpr, WEIGHT_SPLIT: tl.constexpr +): + """`1 + w` for a tile's elements and for their rotation partners; heads at or past + WEIGHT_SPLIT (when > 0) take the second weight (the k heads).""" + HALF: tl.constexpr = ROT // 2 + d = tl.arange(0, D) + rot = d < ROT + partner = tl.where(d < HALF, d + HALF, tl.where(rot, d - HALF, d)) + w = tl.load(weight + d).to(tl.float32) + wp = tl.load(weight + partner).to(tl.float32) + if WEIGHT_SPLIT > 0: + is_k = (heads >= WEIGHT_SPLIT)[:, None] + wk = tl.load(weight_k + d).to(tl.float32) + wkp = tl.load(weight_k + partner).to(tl.float32) + w2 = tl.where(is_k, wk[None, :], w[None, :]) + 1.0 + wp2 = tl.where(is_k, wkp[None, :], wp[None, :]) + 1.0 + else: + w2 = w[None, :] + 1.0 + wp2 = wp[None, :] + 1.0 + return w2, wp2 + + +@triton.jit +def _norm_rope_apply(src, dst, heads, head_mask, w2, wp2, c, s, eps, D: tl.constexpr, ROT: tl.constexpr, + dst_head_stride): # fmt: skip + """`qwen4_exp_qk_norm_rope`'s tile with its static factors given: rms over each head, + `* (1 + w)`, rotate; fp32 throughout, one rounding at the store. Returns the stored + (rounded) tile.""" + HALF: tl.constexpr = ROT // 2 + d = tl.arange(0, D) + rot = d < ROT + partner = tl.where(d < HALF, d + HALF, tl.where(rot, d - HALF, d)) + hm = head_mask[:, None] + x = tl.load(src + heads[:, None] * D + d[None, :], mask=hm, other=0.0).to(tl.float32) + xp = tl.load(src + heads[:, None] * D + partner[None, :], mask=hm & rot[None, :], other=0.0).to( + tl.float32 + ) + rstd = tl.rsqrt(tl.sum(x * x, axis=1) / D + eps) + y = x * rstd[:, None] * w2 + yp = xp * rstd[:, None] * wp2 + out = (y * c[None, :] + yp * s[None, :]).to(dst.dtype.element_ty) + tl.store(dst + heads[:, None] * dst_head_stride + d[None, :], out, mask=hm) + return out + + +@triton.jit +def _compress_factors(pos_t, pos_h, pos_w, cos_sin, cos_sin_stride, ROTARY_PAIRS: tl.constexpr, + IS_MROPE: tl.constexpr, MROPE_H: tl.constexpr, MROPE_W: tl.constexpr, live): # fmt: skip + """The compressed key's rotation factors (QSA's `_qsa_gemma_norm_rope` reads them at + the group's first-token coordinates).""" + pairs = tl.arange(0, ROTARY_PAIRS) + if IS_MROPE: + height_pair = (pairs % 3 == 1) & (pairs < 3 * MROPE_H) + width_pair = (pairs % 3 == 2) & (pairs < 3 * MROPE_W) + position = tl.where(height_pair, pos_h, tl.where(width_pair, pos_w, pos_t)) + else: + position = pos_t + cosine = tl.load(cos_sin + position * cos_sin_stride + pairs, mask=live, other=0.0) + sine = tl.load(cos_sin + position * cos_sin_stride + ROTARY_PAIRS + pairs, mask=live, other=0.0) + return cosine, sine + + +@triton.jit +def _compress_norm_rope( + x, weight, cosine, sine, eps, HEAD_DIM: tl.constexpr, ROTARY_PAIRS: tl.constexpr +): + """QSA's `_qsa_gemma_norm_rope` for one row with its static factors given (`weight` is + `1 + w`; the fp32 table makes its intermediate rounding to the table's dtype a no-op).""" + x_f32 = x.to(tl.float32) + reciprocal_rms = tl.rsqrt(tl.sum(x_f32 * x_f32, axis=1) / HEAD_DIM + eps) + normalized = x_f32 * reciprocal_rms[:, None] * weight[None, :] + rotated, passthrough = tl.split( + tl.permute(tl.reshape(normalized, (1, 2, HEAD_DIM // 2)), (0, 2, 1)) + ) + first, second = tl.split(tl.permute(tl.reshape(rotated, (1, 2, ROTARY_PAIRS)), (0, 2, 1))) + output_first = first * cosine[None, :] - second * sine[None, :] + output_second = second * cosine[None, :] + first * sine[None, :] + rotated = tl.reshape( + tl.permute(tl.join(output_first, output_second), (0, 2, 1)), (1, HEAD_DIM // 2) + ) + return tl.reshape( + tl.permute(tl.join(rotated, passthrough), (0, 2, 1)), + (1, HEAD_DIM), + ).to(tl.bfloat16) + + +@triton.jit(do_not_specialize=["num_tokens", "num_requests", "table_width", "num_pages"]) +def _attn_decode_prologue_kernel( + proj, + stride_proj, + index_q, + stride_iq_token, + stride_iq_head, + positions, + num_tokens, + cos_sin, + stride_cs, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + eps, + k_cache, + v_cache, + s_kc_page, + s_kc_head, + s_kc_tok, + s_vc_page, + s_vc_head, + s_vc_tok, + request_indices, + logical_positions, + block_table, + s_bt_req, + s_bt_col, + num_requests, + table_width, + num_pages, + index_cache, + s_ic_page, + s_ic_slot, + position_cache, + s_pc_page, + s_pc_slot, + s_pc_axis, + coords, + s_c_row, + s_c_axis, + HQ: tl.constexpr, + HK: tl.constexpr, + D: tl.constexpr, + IH: tl.constexpr, + ID: tl.constexpr, + IQ_COL: tl.constexpr, + ROT: tl.constexpr, + S1: tl.constexpr, + S2: tl.constexpr, + IS1: tl.constexpr, + IS2: tl.constexpr, + HPB: tl.constexpr, + MAIN_PROGRAMS: tl.constexpr, + TPB: tl.constexpr, + CR: tl.constexpr, + ROTARY_PAIRS: tl.constexpr, + IS_MROPE: tl.constexpr, + MROPE_H: tl.constexpr, + MROPE_W: tl.constexpr, + USE_PDL: tl.constexpr, +): + token = tl.program_id(0).to(tl.int64) + program = tl.program_id(1) + if USE_PDL: + # The index scores read nothing of this launch before their own wait. + tl.extra.cuda.gdc_launch_dependents() + row = proj + token * stride_proj + + # The row's page and slot (the K/V pool, the index plane and the coordinates share + # the table); an unmapped row writes no cache. Host-filled: read before the wait. + request = tl.load(request_indices + token).to(tl.int64) + logical = tl.load(logical_positions + token).to(tl.int64) + column = logical // TPB + slot = logical % TPB + live = (request >= 0) & (request < num_requests) & (logical >= 0) & (column < table_width) + page = tl.load(block_table + request * s_bt_req + column * s_bt_col, mask=live, other=-1).to( + tl.int64 + ) + live = live & (page >= 0) & (page < num_pages) + + if program < MAIN_PROGRAMS: + main_heads = program * HPB + tl.arange(0, HPB) + c, s = _rope_factors(positions, num_tokens, token, cos_sin, stride_cs, D, ROT, S1, S2) + w2, wp2 = _norm_weights(main_heads, q_weight, k_weight, D, ROT, HQ) + if USE_PDL: + tl.extra.cuda.gdc_wait() + out = _norm_rope_apply( + row, row, main_heads, main_heads < HQ + HK, w2, wp2, c, s, eps, D, ROT, D + ) + if program * HPB + HPB > HQ: + d = tl.arange(0, D) + is_k = (main_heads >= HQ) & (main_heads < HQ + HK) & live + tl.store( + k_cache + + page * s_kc_page + + (main_heads - HQ).to(tl.int64)[:, None] * s_kc_head + + slot * s_kc_tok + + d[None, :], + out.to(k_cache.dtype.element_ty), + mask=is_k[:, None], + ) + elif program == MAIN_PROGRAMS: + index_heads = tl.arange(0, IH) + c, s = _rope_factors(positions, num_tokens, token, cos_sin, stride_cs, ID, ROT, IS1, IS2) + w2, wp2 = _norm_weights(index_heads, index_q_weight, index_q_weight, ID, ROT, 0) + if USE_PDL: + tl.extra.cuda.gdc_wait() + _norm_rope_apply( + row + IQ_COL, index_q + token * stride_iq_token, index_heads, index_heads < IH, + w2, wp2, c, s, eps, ID, ROT, stride_iq_head, + ) # fmt: skip + else: + # Before the wait: the coordinates; at a group's last token, the group's earlier + # raw keys (written by earlier steps), the first token's coordinates, their + # rotation and the key norm weight. + axes = tl.arange(0, 4) + coordinate = tl.load(coords + token * s_c_row + axes * s_c_axis, mask=axes < 3, other=0) + dims = tl.arange(0, ID) + group_end = live & ((logical + 1) % CR == 0) + offsets = tl.arange(0, CR) + members = logical - (CR - 1) + offsets + member_pages = tl.load( + block_table + request * s_bt_req + (tl.maximum(members, 0) // TPB) * s_bt_col, + mask=group_end, + other=0, + ).to(tl.int64) + group = tl.load( + index_cache + + member_pages[:, None] * s_ic_page + + (tl.maximum(members, 0) % TPB)[:, None] * s_ic_slot + + dims[None, :], + mask=group_end & (offsets < CR - 1)[:, None], + other=0.0, + ) + first = tl.maximum(logical - (CR - 1), 0) + first_page = tl.load( + block_table + request * s_bt_req + (first // TPB) * s_bt_col, mask=group_end, other=0 + ).to(tl.int64) + first_base = position_cache + first_page * s_pc_page + (first % TPB) * s_pc_slot + pos_t = tl.load(first_base, mask=group_end, other=0) + pos_h = tl.load(first_base + s_pc_axis, mask=group_end, other=0) + pos_w = tl.load(first_base + 2 * s_pc_axis, mask=group_end, other=0) + cosine, sine = _compress_factors( + pos_t, + pos_h, + pos_w, + cos_sin, + stride_cs, + ROTARY_PAIRS, + IS_MROPE, + MROPE_H, + MROPE_W, + group_end, + ) + key_weight = tl.load(index_k_weight + dims).to(tl.float32) + 1.0 + if USE_PDL: + tl.extra.cuda.gdc_wait() + if live: + # v, as the K/V append stores it. + d = tl.arange(0, D) + for h in tl.static_range(HK): + v = tl.load(row + (HQ + HK + h) * D + d) + tl.store( + v_cache + page * s_vc_page + h * s_vc_head + slot * s_vc_tok + d, + v.to(v_cache.dtype.element_ty), + ) + # The coordinates, as the index store writes them. + tl.store( + position_cache + page * s_pc_page + slot * s_pc_slot + axes * s_pc_axis, + coordinate.to(position_cache.dtype.element_ty), + mask=axes < 3, + ) + # The raw index key -- or, at a group's last token, the group compressed in + # its place (the compression kernel overwrites the same slot). + current_k = tl.load(row + IQ_COL + IH * ID + dims) + stored_k = current_k.to(index_cache.dtype.element_ty) + if (logical + 1) % CR == 0: + group = tl.where(offsets[:, None] == CR - 1, stored_k[None, :], group) + pooled = (tl.sum(group.to(tl.float32), axis=0) / CR).to(tl.bfloat16) + stored_k = tl.reshape( + _compress_norm_rope( + pooled[None, :], key_weight, cosine, sine, eps, ID, ROTARY_PAIRS + ), + (ID,), + ) + tl.store(index_cache + page * s_ic_page + slot * s_ic_slot + dims, stored_k) + + +def _pdl() -> bool: + return os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" + + +class Qwen4ExpAttnDecodePrologue(OpWrapper): + """`qwen4_exp_qk_norm_rope`, `qsa_paged_kv_store`, `qsa_paged_index_store` and + `qsa_prefill_compress` of one sparse decode attention layer, in one launch. + + `proj` `[T, W]` is the layer's fused input projection: q heads `[0, Hq)`, k heads + `[Hq, Hq + Hk)` and v heads `[Hq + Hk, Hq + 2 Hk)` of `head_dim`, the index query + heads at `index_q_col` and the raw index key right after them. Writes, as the four + entries do: q and k normalized and rotated in place on `proj`; the index query into + `index_q`; k and v into the paged K/V pool and the raw index key (or, at a group's + last token, the compressed group) and the coordinates into the index plane and the + position cache, at the row's page and slot. Returns `(proj, index_q)`. + + Fusion boundary: as the four entries; the projection before and the index scores + after are the caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="decode, 1 row at position 8699 (a group's last token)", spec=dict(pos=(8699,))), + Cell(why="decode, 1 row at position 8700 (no compression)", spec=dict(pos=(8700,))), + Cell( + why="decode, 4 rows, two group ends, one page boundary", + spec=dict(pos=(8699, 8700, 4095, 31)), + ), + Cell( + why="decode, 2 rows, one unmapped (request out of range)", + spec=dict(pos=(8699, 1201), unmapped=1), + ), + ) + + note = """ + Metadata consumed: the QSA request indices, logical positions and block table; the + position ids as `[3, T]` (the rope) and as `[T, 3]` (the coordinates stored). + Preconditions: those of the four entries; the index plane one KV head of + `index_head_dim` (bf16), the position cache int32 `[pages, tpb, >=3]`; decode rows, + one per request (a group's earlier tokens are already in the index plane). A row + whose page lookup fails writes no cache, as the stores and compression skip it. + + Numerics: the same arithmetic as the four entries; outputs bitwise theirs. + """ + + def __call__( + self, + proj: torch.Tensor, + index_q: torch.Tensor, + position_ids: torch.Tensor, + coords: torch.Tensor, + cos_sin: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + index_cache: torch.Tensor, + position_cache: torch.Tensor, + request_indices: torch.Tensor, + logical_positions: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + compress_ratio: int, + num_heads_q: int, + num_heads_k: int, + head_dim: int, + index_heads: int, + index_head_dim: int, + index_q_col: int, + rotary_dim: int, + eps: float, + mrope_section1: int, + mrope_section2: int, + index_mrope_section1: int = 0, + index_mrope_section2: int = 0, + compress_mrope_section: Optional[Tuple[int, int, int]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + num_tokens = proj.shape[0] + if num_tokens == 0: + return proj, index_q + hpb, warps = _qk_launch(num_tokens) + main_programs = triton.cdiv(num_heads_q + num_heads_k, hpb) + section = compress_mrope_section or (0, 0, 0) + pdl = _pdl() + _attn_decode_prologue_kernel[(num_tokens, main_programs + 2)]( + proj, + proj.stride(0), + index_q, + index_q.stride(0), + index_q.stride(1), + position_ids, + num_tokens, + cos_sin, + cos_sin.stride(0), + q_weight, + k_weight, + index_q_weight, + index_k_weight, + eps, + k_cache, + v_cache, + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(2), + v_cache.stride(0), + v_cache.stride(1), + v_cache.stride(2), + request_indices, + logical_positions, + block_table, + block_table.stride(0), + block_table.stride(1), + block_table.shape[0], + block_table.shape[1], + k_cache.shape[0], + index_cache, + index_cache.stride(0), + index_cache.stride(1), + position_cache, + position_cache.stride(0), + position_cache.stride(1), + position_cache.stride(2), + coords, + coords.stride(0), + coords.stride(1), + HQ=num_heads_q, + HK=num_heads_k, + D=head_dim, + IH=index_heads, + ID=index_head_dim, + IQ_COL=index_q_col, + ROT=rotary_dim, + S1=mrope_section1, + S2=mrope_section2, + IS1=index_mrope_section1, + IS2=index_mrope_section2, + HPB=hpb, + MAIN_PROGRAMS=main_programs, + TPB=tokens_per_block, + CR=compress_ratio, + ROTARY_PAIRS=cos_sin.shape[1] // 2, + IS_MROPE=compress_mrope_section is not None, + MROPE_H=section[1], + MROPE_W=section[2], + USE_PDL=pdl, + num_warps=warps, + launch_pdl=pdl, + ) + return proj, index_q + + def reference( + self, + proj: torch.Tensor, + index_q: torch.Tensor, + position_ids: torch.Tensor, + coords: torch.Tensor, + cos_sin: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + index_cache: torch.Tensor, + position_cache: torch.Tensor, + request_indices: torch.Tensor, + logical_positions: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + compress_ratio: int, + num_heads_q: int, + num_heads_k: int, + head_dim: int, + index_heads: int, + index_head_dim: int, + index_q_col: int, + rotary_dim: int, + eps: float, + mrope_section1: int, + mrope_section2: int, + index_mrope_section1: int = 0, + index_mrope_section2: int = 0, + compress_mrope_section: Optional[Tuple[int, int, int]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """The four entries in their composition order, on the given buffers.""" + rows = proj.shape[0] + _qk_norm_rope( + proj, index_q, position_ids, cos_sin, q_weight, k_weight, index_q_weight, + num_heads_q=num_heads_q, num_heads_k=num_heads_k, head_dim=head_dim, + index_heads=index_heads, index_head_dim=index_head_dim, index_q_col=index_q_col, + rotary_dim=rotary_dim, eps=eps, mrope_section1=mrope_section1, + mrope_section2=mrope_section2, index_mrope_section1=index_mrope_section1, + index_mrope_section2=index_mrope_section2, + ) # fmt: skip + heads = proj.view(rows, -1, head_dim) + k = heads[:, num_heads_q : num_heads_q + num_heads_k] + v = heads[:, num_heads_q + num_heads_k : num_heads_q + 2 * num_heads_k] + _kv_store( + k, + v, + k_cache, + v_cache, + request_indices, + logical_positions, + block_table, + tokens_per_block, + ) + token_k = proj[:, index_q_col + index_heads * index_head_dim :][:, :index_head_dim].reshape( + rows, 1, index_head_dim + ) + _index_store( + token_k, coords, index_cache, position_cache, request_indices, logical_positions, + block_table, tokens_per_block, + ) # fmt: skip + _compress( + logical_positions=logical_positions, request_indices=request_indices, + block_table=block_table, index_cache=index_cache, position_cache=position_cache, + k_norm_weight=index_k_weight, cos_sin=cos_sin, eps=eps, + tokens_per_block=tokens_per_block, compress_ratio=compress_ratio, + mrope_section=compress_mrope_section, + ) # fmt: skip + return proj, index_q + + def is_valid(self, proj, index_q, position_ids, coords, cos_sin, *args, **kwargs) -> None: + assert proj.stride(1) == 1 and proj.dtype == torch.bfloat16, "bf16 proj, unit column stride" + assert position_ids.is_contiguous() and position_ids.dtype == torch.int32, ( + "int32 [3, T] positions" + ) + assert coords.shape == (proj.shape[0], 3), "coords [T, 3]" + assert cos_sin.dtype == torch.float32, ( + "an fp32 cos|sin table (the compressed key's rotation)" + ) + + +qwen4_exp_attn_decode_prologue = Qwen4ExpAttnDecodePrologue() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_qk_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_qk_norm_rope.py new file mode 100644 index 000000000000..cb1a60d7c181 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_qk_norm_rope.py @@ -0,0 +1,1201 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Gemma QK norm + partial NeoX rope for both of an attention layer's planes, one launch. + +The main plane (q and k heads of the fused input projection) is rewritten in +place; the QSA index plane's query heads are read from the same projection row +and written, normalized and rotated, into a separate buffer. The rotation reads +the target's own fp32 `[positions, rotary_dim]` cos|sin table. + +Two kernels, one result: a Triton kernel, and from `_CUTE_MIN_TOKENS` tokens (a mixed +step's prefill chunk) a persistent CuTe DSL kernel that is bitwise the Triton kernel's +eight-head launch and streams the projection at about twice its rate. +""" + +from __future__ import annotations + +import struct +from typing import Dict, Tuple + +import cutlass +import cutlass.cute as cute +import torch +import triton +import triton.language as tl +from cutlass._mlir.dialects import llvm +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.cute_dsl_kernels.blackwell.utils import make_ptr + + +@triton.jit +def _norm_rope_tile( + src, + dst, + heads, + head_mask, + weight, + positions, + num_tokens, + token, + cos_sin, + stride_cs, + eps, + D: tl.constexpr, + ROT: tl.constexpr, + S1: tl.constexpr, + S2: tl.constexpr, + WEIGHT_SPLIT: tl.constexpr, + weight_k, + dst_head_stride, +): + """One `[len(heads), D]` tile: rms over each head, `* (1 + w)`, rotate the + first ROT dims in NeoX halves. fp32 throughout, one rounding at the store.""" + HALF: tl.constexpr = ROT // 2 + d = tl.arange(0, D) + rot = d < ROT + partner = tl.where(d < HALF, d + HALF, tl.where(rot, d - HALF, d)) + hm = head_mask[:, None] + x = tl.load(src + heads[:, None] * D + d[None, :], mask=hm, other=0.0).to(tl.float32) + xp = tl.load(src + heads[:, None] * D + partner[None, :], mask=hm & rot[None, :], other=0.0).to( + tl.float32 + ) + w = tl.load(weight + d).to(tl.float32) + wp = tl.load(weight + partner).to(tl.float32) + if WEIGHT_SPLIT > 0: + # Heads at or past WEIGHT_SPLIT take the second weight (the k heads). + is_k = (heads >= WEIGHT_SPLIT)[:, None] + wk = tl.load(weight_k + d).to(tl.float32) + wkp = tl.load(weight_k + partner).to(tl.float32) + w2 = tl.where(is_k, wk[None, :], w[None, :]) + 1.0 + wp2 = tl.where(is_k, wkp[None, :], wp[None, :]) + 1.0 + else: + w2 = w[None, :] + 1.0 + wp2 = wp[None, :] + 1.0 + rstd = tl.rsqrt(tl.sum(x * x, axis=1) / D + eps) + y = x * rstd[:, None] * w2 + yp = xp * rstd[:, None] * wp2 + # Interleaved mRoPE: pair j reads coordinate 1 when j % 3 == 1 and + # j < 3 * S1, coordinate 2 when j % 3 == 2 and j < 3 * S2, else 0. + j = d % HALF + p0 = tl.load(positions + token) + p1 = tl.load(positions + num_tokens + token) + p2 = tl.load(positions + 2 * num_tokens + token) + pos = tl.where((j % 3 == 1) & (j < 3 * S1), p1, tl.where((j % 3 == 2) & (j < 3 * S2), p2, p0)) + row = cos_sin + pos.to(tl.int64) * stride_cs + c = tl.load(row + j, mask=rot, other=1.0) + s = tl.load(row + HALF + j, mask=rot, other=0.0) + s = tl.where(d < HALF, -s, s) + out = y * c[None, :] + yp * s[None, :] + tl.store( + dst + heads[:, None] * dst_head_stride + d[None, :], + out.to(dst.dtype.element_ty), + mask=hm, + ) + + +@triton.jit(do_not_specialize=["num_tokens"]) +def _qwen4_exp_qk_norm_rope_kernel( + proj, + stride_proj, + index_q, + stride_iq_token, + stride_iq_head, + positions, + num_tokens, + cos_sin, + stride_cs, + q_weight, + k_weight, + index_q_weight, + eps, + HQ: tl.constexpr, + HK: tl.constexpr, + D: tl.constexpr, + IH: tl.constexpr, + ID: tl.constexpr, + IQ_COL: tl.constexpr, + ROT: tl.constexpr, + S1: tl.constexpr, + S2: tl.constexpr, + IS1: tl.constexpr, + IS2: tl.constexpr, + HPB: tl.constexpr, + MAIN_PROGRAMS: tl.constexpr, +): + token = tl.program_id(0).to(tl.int64) + program = tl.program_id(1) + row = proj + token * stride_proj + if program < MAIN_PROGRAMS: + main_heads = program * HPB + tl.arange(0, HPB) + _norm_rope_tile( + row, row, main_heads, main_heads < HQ + HK, q_weight, positions, num_tokens, token, + cos_sin, stride_cs, eps, D, ROT, S1, S2, HQ, k_weight, D, + ) # fmt: skip + else: + index_heads = tl.arange(0, IH) + _norm_rope_tile( + row + IQ_COL, index_q + token * stride_iq_token, index_heads, index_heads < IH, + index_q_weight, positions, num_tokens, token, cos_sin, stride_cs, eps, + ID, ROT, IS1, IS2, 0, index_q_weight, stride_iq_head, + ) # fmt: skip + + +def _launch(num_tokens: int) -> Tuple[int, int]: + """Main-plane heads per program and warps per program, by token count. + + Measured on GB300 (microbench `qk-norm-rope-decode`, CUDA-graph replay, 12 + rotating layers, us per call; HPB/warps): + + tokens 1/2 2/2 2/4 4/4 8/8 32/8 + 1 1.76 1.84 1.87 1.99 1.92 2.82 + 8 1.83 1.97 1.97 2.02 2.05 2.91 + 32 2.11 2.14 2.18 2.14 2.16 2.99 + 64 2.48 2.39 2.47 2.44 2.48 3.03 + 128 3.24 3.19 3.46 3.12 2.96 3.12 + 2048 29.7 25.4 30.8 24.0 22.6 25.8 + 8192 115 98.5 121 95.7 89.5 122 + """ + if num_tokens <= 16: + return 1, 2 + if num_tokens <= 64: + return 2, 2 + return 8, 8 + + +#: This checkpoint's geometry: 24 q + 2 k heads of 256 at the head of the +#: 14080-wide fused projection row, 4 index q heads of 128 at column 13312, a +#: 64-dim rotary prefix, `mrope_section` [11, 11, 10] on the main plane and +#: coordinate 0 alone on the index plane. +CKPT = dict( + num_heads_q=24, + num_heads_k=2, + head_dim=256, + index_heads=4, + index_head_dim=128, + index_q_col=13312, + proj_width=14080, + index_scratch_heads=6, + rotary_dim=64, + mrope_section1=11, + mrope_section2=10, + index_mrope_section1=0, + index_mrope_section2=0, +) + + +# ---- CuTe DSL path: prefill-size token counts ------------------------------------------- + +#: Token count from which the op runs the CuTe kernel. GB300, graph replay, us per call (Triton +#: -> CuTe): 512 8.4 -> 10.3, 768 12.3 -> 8.2, 1024 14.5 -> 8.3, 4096 45.2 -> 18.5, 16384 +#: 182 -> 82.7. The Triton kernel runs its eight-head launch from 65 tokens on, and the CuTe +#: kernel is bitwise that launch. +_CUTE_MIN_TOKENS = 1024 +#: A token's stage in a warp's ring: the lanes' four chunks (32 x 16 bytes each), the cos and +#: sin of the 32 rotary pairs, and three positions. +_STAGE = 4 * 32 * 16 + 2 * 32 * 4 + 16 +_STAGE_CS = 4 * 32 * 16 +_STAGE_POS = _STAGE_CS + 2 * 32 * 4 +#: Stages in a warp's ring (the next token's loads in flight while one is formed). +_RING = 2 +_CTAS_PER_SM = 3 + + +def _iv(x, loc, ip): + return x.ir_value(loc=loc, ip=ip) + + +def _asm(ret, args, text, constraints, loc, ip, side_effects=False): + return llvm.inline_asm( + ret, + [_iv(a, loc, ip) for a in args], + text, + constraints, + has_side_effects=side_effects, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _asm4(args, text, constraints, loc, ip, side_effects=False): + ty = T.i32() + res = _asm( + llvm.StructType.get_literal([ty] * 4), args, text, constraints, loc, ip, side_effects + ) + return [cutlass.Int32(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip)) for i in range(4)] + + +@dsl_user_op +def _ldg_v4(addr, *, loc=None, ip=None): + return _asm4( + [cutlass.Int64(addr)], + "ld.global.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _ldg_b32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm(T.i32(), [cutlass.Int64(addr)], "ld.global.b32 $0, [$1];", "=r,l", loc, ip, True) + ) + + +@dsl_user_op +def _stg_v4(addr, off: int, a, b, c, d, *, loc=None, ip=None): + """16 bytes at `addr + off`.""" + _asm( + T.i32(), + [cutlass.Int64(addr)] + [cutlass.Int32(v) for v in (a, b, c, d)], + "{ st.global.v4.b32 [$1+%d], {$2, $3, $4, $5}; mov.u32 $0, 0; }" % off, + "=r,l,r,r,r,r", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _lds_v4(addr, *, loc=None, ip=None): + return _asm4( + [cutlass.Int32(addr)], + "ld.shared.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,r", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _lds_b32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm(T.i32(), [cutlass.Int32(addr)], "ld.shared.b32 $0, [$1];", "=r,r", loc, ip, True) + ) + + +@dsl_user_op +def _cp16(dst, src, *, loc=None, ip=None): + _asm( + T.i32(), + [cutlass.Int32(dst), cutlass.Int64(src)], + "{ cp.async.cg.shared.global [$1], [$2], 16; mov.u32 $0, 0; }", + "=r,r,l", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _cp4(dst, src, pred, *, loc=None, ip=None): + """4 bytes where `pred` is not 0.""" + _asm( + T.i32(), + [cutlass.Int32(dst), cutlass.Int64(src), cutlass.Int32(pred)], + "{ .reg .pred p; setp.ne.b32 p, $3, 0; @p cp.async.ca.shared.global [$1], [$2], 4; " + "mov.u32 $0, 0; }", + "=r,r,l,r", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _cp_commit(*, loc=None, ip=None): + _asm(T.i32(), [], "{ cp.async.commit_group; mov.u32 $0, 0; }", "=r", loc, ip, True) + + +@dsl_user_op +def _cp_wait(n: int, *, loc=None, ip=None): + _asm(T.i32(), [], "{ cp.async.wait_group %d; mov.u32 $0, 0; }" % n, "=r", loc, ip, True) + + +@dsl_user_op +def _bf16_lo(w, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Int32(w)], + "{ .reg .b32 t; shl.b32 t, $1, 16; mov.b32 $0, t; }", + "=f,r", + loc, + ip, + ) + ) + + +@dsl_user_op +def _bf16_hi(w, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Int32(w)], + "{ .reg .b32 t; and.b32 t, $1, 0xFFFF0000; mov.b32 $0, t; }", + "=f,r", + loc, + ip, + ) + ) + + +@dsl_user_op +def _unpack2(w, *, loc=None, ip=None): + """A bf16x2 word -> its two values as an fp32x2 pair, the low element first.""" + return cutlass.Int64( + _asm( + T.i64(), + [cutlass.Int32(w)], + "{ .reg .b32 a, b; shl.b32 a, $1, 16; and.b32 b, $1, 0xFFFF0000; mov.b64 $0, {a, b}; }", + "=l,r", + loc, + ip, + ) + ) + + +@dsl_user_op +def _pair(lo, hi, *, loc=None, ip=None): + """Two fp32 bit patterns -> one fp32x2 pair.""" + return cutlass.Int64( + _asm( + T.i64(), + [cutlass.Int32(lo), cutlass.Int32(hi)], + "mov.b64 $0, {$1, $2};", + "=l,r,r", + loc, + ip, + ) + ) + + +@dsl_user_op +def _splat2(x, *, loc=None, ip=None): + return cutlass.Int64( + _asm(T.i64(), [cutlass.Float32(x)], "mov.b64 $0, {$1, $1};", "=l,f", loc, ip) + ) + + +@dsl_user_op +def _mul2(a, b, *, loc=None, ip=None): + return cutlass.Int64( + _asm( + T.i64(), + [cutlass.Int64(a), cutlass.Int64(b)], + "mul.rn.f32x2 $0, $1, $2;", + "=l,l,l", + loc, + ip, + ) + ) + + +@dsl_user_op +def _fma2(a, b, c, *, loc=None, ip=None): + return cutlass.Int64( + _asm( + T.i64(), + [cutlass.Int64(a), cutlass.Int64(b), cutlass.Int64(c)], + "fma.rn.f32x2 $0, $1, $2, $3;", + "=l,l,l,l", + loc, + ip, + ) + ) + + +@dsl_user_op +def _add2(a, b, *, loc=None, ip=None): + return cutlass.Int64( + _asm( + T.i64(), + [cutlass.Int64(a), cutlass.Int64(b)], + "add.rn.f32x2 $0, $1, $2;", + "=l,l,l", + loc, + ip, + ) + ) + + +@dsl_user_op +def _hsum2(a, *, loc=None, ip=None): + """The low plus the high element of a pair.""" + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Int64(a)], + "{ .reg .f32 x, y; mov.b64 {x, y}, $1; add.rn.f32 $0, x, y; }", + "=f,l", + loc, + ip, + ) + ) + + +@dsl_user_op +def _bf16x2_of(a, *, loc=None, ip=None): + """A pair rounded (RN) into one bf16x2 word, the low element low.""" + return cutlass.Int32( + _asm( + T.i32(), + [cutlass.Int64(a)], + "{ .reg .f32 x, y; mov.b64 {x, y}, $1; cvt.rn.bf16x2.f32 $0, y, x; }", + "=r,l", + loc, + ip, + ) + ) + + +@dsl_user_op +def _neg_bits(w, *, loc=None, ip=None): + """An fp32 bit pattern with its sign flipped (the Triton kernel's `xor`).""" + return cutlass.Int32( + _asm(T.i32(), [cutlass.Int32(w)], "xor.b32 $0, $1, 0x80000000;", "=r,r", loc, ip) + ) + + +@dsl_user_op +def _fma(a, b, c, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _mul(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _add(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "add.rn.f32 $0, $1, $2;", + "=f,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _rstd(ss, eps, width_hex: str, *, loc=None, ip=None): + """`rsqrt(ss / width + eps)` as the Triton kernel's `div.full`, `add`, `rsqrt.approx.ftz`: + the `add` has no rounding mode there, so ptxas fuses it with the division (a multiply by + the reciprocal) into one FFMA -- here the same.""" + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Float32(ss), cutlass.Float32(eps)], + "{ .reg .f32 t; div.full.f32 t, $1, " + width_hex + "; add.f32 t, t, $2; " + "rsqrt.approx.ftz.f32 $0, t; }", + "=f,f,f", + loc, + ip, + ) + ) + + +def _hex(x: float) -> str: + return "0f%08X" % struct.unpack(" bool: + """The geometry and layouts the CuTe kernel covers; anything else runs the Triton kernel.""" + bf = torch.bfloat16 + return ( + num_heads_q % 4 == 0 + and num_heads_k == 2 + and head_dim == 256 + and index_heads == 4 + and index_head_dim == 128 + and rotary_dim == 64 + and index_q_col % 8 == 0 + and proj.dtype == index_q.dtype == bf + and all(w.dtype == bf and w.is_contiguous() and w.data_ptr() % 16 == 0 for w in weights) + and cos_sin.dtype == torch.float32 + and cos_sin.stride(1) == 1 + and proj.data_ptr() % 16 == 0 + and proj.stride(0) % 8 == 0 + and index_q.data_ptr() % 16 == 0 + and index_q.stride(0) % 8 == 0 + and index_q.stride(1) == index_head_dim + ) + + +def _launch_cute( + proj: torch.Tensor, + index_q: torch.Tensor, + position_ids: torch.Tensor, + cos_sin: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + num_heads_q: int, + index_q_col: int, + eps: float, + sections: Tuple[int, int, int, int], +) -> None: + key = (num_heads_q, index_q_col) + tuple(sections) + fn = _RING_KERNELS.get(key) + if fn is None: + gm = cute.AddressSpace.gmem + bf = cutlass.BFloat16 + fn = cute.compile( + _NormRopeRing(num_heads_q, index_q_col, *sections), + from_dlpack(proj, assumed_align=16).mark_layout_dynamic(leading_dim=1), + make_ptr(bf, index_q.data_ptr(), gm, assumed_align=16), + make_ptr(cutlass.Int32, position_ids.data_ptr(), gm, assumed_align=4), + make_ptr(cutlass.Float32, cos_sin.data_ptr(), gm, assumed_align=4), + make_ptr(bf, q_weight.data_ptr(), gm, assumed_align=16), + make_ptr(bf, k_weight.data_ptr(), gm, assumed_align=16), + make_ptr(bf, index_q_weight.data_ptr(), gm, assumed_align=16), + *(cutlass.Int32(1),) * 4, + cutlass.Float32(1e-6), + cutlass.Int32(1), + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--opt-level 2 --enable-tvm-ffi", + ) + _RING_KERNELS[key] = fn + dev = proj.device.index + if dev not in _SM_COUNT: + _SM_COUNT[dev] = torch.cuda.get_device_properties(proj.device).multi_processor_count + tokens = proj.shape[0] + fn( + proj, + index_q.data_ptr(), + position_ids.data_ptr(), + cos_sin.data_ptr(), + q_weight.data_ptr(), + k_weight.data_ptr(), + index_q_weight.data_ptr(), + tokens, + proj.stride(0), + index_q.stride(0), + cos_sin.stride(0), + eps, + min(tokens, _SM_COUNT[dev] * _CTAS_PER_SM), + ) + + +def _launch_triton( + proj: torch.Tensor, + index_q: torch.Tensor, + position_ids: torch.Tensor, + cos_sin: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + num_heads_q: int, + num_heads_k: int, + head_dim: int, + index_heads: int, + index_head_dim: int, + index_q_col: int, + rotary_dim: int, + eps: float, + sections: Tuple[int, int, int, int], +) -> None: + num_tokens = proj.shape[0] + hpb, warps = _launch(num_tokens) + main_programs = triton.cdiv(num_heads_q + num_heads_k, hpb) + _qwen4_exp_qk_norm_rope_kernel[(num_tokens, main_programs + 1)]( + proj, + proj.stride(0), + index_q, + index_q.stride(0), + index_q.stride(1), + position_ids, + num_tokens, + cos_sin, + cos_sin.stride(0), + q_weight, + k_weight, + index_q_weight, + eps, + HQ=num_heads_q, + HK=num_heads_k, + D=head_dim, + IH=index_heads, + ID=index_head_dim, + IQ_COL=index_q_col, + ROT=rotary_dim, + S1=sections[0], + S2=sections[1], + IS1=sections[2], + IS2=sections[3], + HPB=hpb, + MAIN_PROGRAMS=main_programs, + num_warps=warps, + ) + + +class Qwen4ExpQkNormRope(OpWrapper): + """Both planes of one attention layer's QK norm + rope, in one launch. + + Main plane, in place on `proj`: heads `[0, Hq)` (weight `q_weight`) and + `[Hq, Hq + Hk)` (weight `k_weight`), each `head_dim` wide at the start of + every row, become `rope(rmsnorm(x) * (1 + w))`. Index plane: the + `index_heads` heads of `index_head_dim` starting at column `index_q_col` of + `proj` are read, normalized with `index_q_weight` and rotated, and written + to `index_q` (`[T, index_heads, index_head_dim]`, any token/head stride, + unit dim stride); `proj`'s index columns are not written. Every other + column of `proj` is left alone. Returns `(proj, index_q)`. + + The rotation is NeoX over the first `rotary_dim` dims of a head: pair `j` + is `(j, j + rotary_dim / 2)`, rotated by `cos_sin[p, j]` and + `cos_sin[p, rotary_dim / 2 + j]` (`cos_sin` is the fp32 `[P, rotary_dim]` + table, cos block then sin block). `position_ids` is the contiguous int32 + `[3, T]`; pair `j` reads coordinate 1 when `j % 3 == 1 and j < 3 * s1`, + coordinate 2 when `j % 3 == 2 and j < 3 * s2`, coordinate 0 otherwise -- + interleaved mRoPE with the plane's own `(s1, s2)`; `(0, 0)` is the plain + rope on coordinate 0. fp32 throughout, one rounding at the store. + + Fusion boundary: norm and rotation of both planes. The projection GEMM + before it, and the KV/index-cache writes and attention after it, are the + caller's. The index plane's raw key is not touched, so the index-cache + store may run before or after this call. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="decode M=1: one token, three distinct mRoPE rows", spec=dict(tokens=1, **CKPT)), + Cell(why="decode M=2", spec=dict(tokens=2, **CKPT)), + Cell(why="decode M=5: odd batch", spec=dict(tokens=5, **CKPT)), + Cell( + why="decode M=16: last token count on the 1-head/2-warp launch", + spec=dict(tokens=16, **CKPT), + ), + Cell( + why="decode M=17: first on the 2-head launch (13 main programs)", + spec=dict(tokens=17, **CKPT), + ), + Cell(why="decode M=33: one past a power of two", spec=dict(tokens=33, **CKPT)), + Cell(why="decode M=64: last on the 2-head launch", spec=dict(tokens=64, **CKPT)), + Cell( + why="decode M=65: first on the 8-head launch (last main program half masked)", + spec=dict(tokens=65, **CKPT), + ), # fmt: skip + Cell(why="decode M=127: odd, below max concurrency", spec=dict(tokens=127, **CKPT)), + Cell(why="decode M=128: max concurrency", spec=dict(tokens=128, **CKPT)), + Cell( + why="decode M=128, positions up to the table end (262144 rows)", + spec=dict(tokens=128, max_position=262144, **CKPT), + ), + Cell(why="prefill 1023 tokens: last on the Triton kernel", spec=dict(tokens=1023, **CKPT)), + Cell(why="prefill 1024 tokens: first on the CuTe kernel", spec=dict(tokens=1024, **CKPT)), + Cell(why="prefill 2048 tokens", spec=dict(tokens=2048, **CKPT)), + Cell(why="prefill 8193 tokens: one past the ISL", spec=dict(tokens=8193, **CKPT)), + Cell(why="prefill 16384 tokens: max_num_tokens", spec=dict(tokens=16384, **CKPT)), + ) + + note = """ + Metadata consumed: none (positions are an operand). Below `_CUTE_MIN_TOKENS` + tokens one Triton kernel, no autotuner: grid `(T, cdiv(Hq + Hk, HPB) + 1)`; + the last program of a token row runs the whole index plane as one + `[index_heads, index_head_dim]` tile. HPB and warps by token count from + `_launch` (its docstring has the table). Replaces, per attention layer, the + CUDA `fused_qk_norm_rope` main call, the index-scratch copy and the second + `fused_qk_norm_rope` call: 6.67 -> 1.76 us at M=1, 9.15 -> 2.96 at M=128, 185 + -> 89.5 at 8192 tokens (graph replay). + + From `_CUTE_MIN_TOKENS` tokens, where the geometry is this checkpoint's (`_cute_ok`: + head dims 256 / 128, 2 k heads, 4 index heads, a 64-dim rotary prefix, 16-byte + aligned rows, contiguous index heads), the CuTe DSL `_NormRopeRing`: three + persistent CTAs an SM, each forming every G-th token, a warp's next token + prefetched by cp.async (its class docstring has the mapping). 16384 tokens: 182 -> 82.7 us in the + microbenchmark (the Triton kernel was L1-bound; the ring runs at ~65% of DRAM peak, + bounded by load latency at three CTAs an SM). + + Numerics: fp32 throughout, one rounding at the store; the CuTe kernel is bitwise + the Triton kernel's eight-head launch (the test checks it at 1024-16384 tokens, + signed zeros included), so the cells' bands cover both. Against the fp64 + reference over the same table every cell is within one output rounding + step (measured max 1.48 ulp of the row scale, relative RMS <= 0.011 ulp). + Against the CUDA pair it replaces, <= 1.41 ulp of the row scale and 0.09 + ulp RMS (that op derives cos/sin with `__sincosf`, whose phase error grows + with position). Three agreeing position rows select the same table row for + every pair, so a text batch is exactly the plain rope on row 0. + + Preconditions in `is_valid`: the main q/k columns end at or before + `index_q_col` (the planes do not overlap, so reading the index plane while + the main plane is rewritten in place is race-free); `index_q` does not + alias `proj`; positions below the table length (an out-of-range position + reads past the table). + + Hard limits: token offsets in int64; positions index the table in int64. + + Not certified: fp16/fp32 activations, `rotary_dim` > either head width, + non-power-of-two head widths or head counts (tl.arange), YaRN, non-Gemma + norm, a k plane on the index side (its scratch slot is not written). + """ + + def __call__( + self, + proj: torch.Tensor, + index_q: torch.Tensor, + position_ids: torch.Tensor, + cos_sin: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + num_heads_q: int, + num_heads_k: int, + head_dim: int, + index_heads: int, + index_head_dim: int, + index_q_col: int, + rotary_dim: int, + eps: float, + mrope_section1: int, + mrope_section2: int, + index_mrope_section1: int = 0, + index_mrope_section2: int = 0, + ) -> Tuple[torch.Tensor, torch.Tensor]: + num_tokens = proj.shape[0] + if num_tokens == 0: + return proj, index_q + sections = (mrope_section1, mrope_section2, index_mrope_section1, index_mrope_section2) + weights = (q_weight, k_weight, index_q_weight) + if num_tokens >= _CUTE_MIN_TOKENS and _cute_ok( + proj, index_q, cos_sin, weights, num_heads_q, num_heads_k, head_dim, index_heads, + index_head_dim, index_q_col, rotary_dim, + ): # fmt: skip + _launch_cute( + proj, index_q, position_ids, cos_sin, *weights, num_heads_q, index_q_col, eps, + sections, + ) # fmt: skip + else: + _launch_triton( + proj, index_q, position_ids, cos_sin, *weights, num_heads_q, num_heads_k, + head_dim, index_heads, index_head_dim, index_q_col, rotary_dim, eps, sections, + ) # fmt: skip + return proj, index_q + + def reference( + self, + proj: torch.Tensor, + index_q: torch.Tensor, + position_ids: torch.Tensor, + cos_sin: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + num_heads_q: int, + num_heads_k: int, + head_dim: int, + index_heads: int, + index_head_dim: int, + index_q_col: int, + rotary_dim: int, + eps: float, + mrope_section1: int, + mrope_section2: int, + index_mrope_section1: int = 0, + index_mrope_section2: int = 0, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """`(proj after, index_q's written values)` in fp64 over the same table, + rounded once; neither input is mutated.""" + num_tokens = proj.shape[0] + half = rotary_dim // 2 + table = cos_sin.double() + pos = position_ids.long() + j = torch.arange(half, device=proj.device) + + def angles(s1: int, s2: int): + coord = torch.zeros(half, dtype=torch.long, device=proj.device) + coord[(j % 3 == 1) & (j < 3 * s1)] = 1 + coord[(j % 3 == 2) & (j < 3 * s2)] = 2 + p = pos[coord, :].t() # [T, half] + return table[p, j][:, None, :], table[p, half + j][:, None, :] + + def plane(x: torch.Tensor, w: torch.Tensor, cos, sin) -> torch.Tensor: + x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * (1.0 + w) + x1, x2 = x[..., :half], x[..., half:rotary_dim] + rotated = torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], -1) + return torch.cat([rotated, x[..., rotary_dim:]], -1) + + out = proj.double().clone() + main_heads = num_heads_q + num_heads_k + main_cols = main_heads * head_dim + w = torch.cat( + [ + q_weight.double().expand(num_heads_q, head_dim), + k_weight.double().expand(num_heads_k, head_dim), + ] + ) + main = out[:, :main_cols].reshape(num_tokens, main_heads, head_dim) + out[:, :main_cols] = plane(main, w, *angles(mrope_section1, mrope_section2)).reshape( + num_tokens, main_cols + ) + iq = proj[:, index_q_col : index_q_col + index_heads * index_head_dim].double() + iq = plane( + iq.reshape(num_tokens, index_heads, index_head_dim), + index_q_weight.double(), + *angles(index_mrope_section1, index_mrope_section2), + ) + return out.to(proj.dtype), iq.to(index_q.dtype) + + def is_valid( + self, + proj: torch.Tensor, + index_q: torch.Tensor, + position_ids: torch.Tensor, + cos_sin: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + num_heads_q: int, + num_heads_k: int, + head_dim: int, + index_heads: int, + index_head_dim: int, + index_q_col: int, + rotary_dim: int, + *unused_args: object, + **unused_kwargs: object, + ) -> None: + num_tokens = proj.shape[0] + assert proj.stride(-1) == 1, "proj must have a unit column stride" + assert (num_heads_q + num_heads_k) * head_dim <= index_q_col, ( + "the index q columns overlap the main q/k heads: the in-place main write races " + "the index read" + ) + assert index_q_col + index_heads * index_head_dim <= proj.shape[1], ( + "the index q columns run past the projection row" + ) + assert index_q.shape == (num_tokens, index_heads, index_head_dim), index_q.shape + assert index_q.stride(-1) == 1, "index_q must have a unit dim stride" + lo, hi = index_q.data_ptr(), proj.data_ptr() + assert lo >= hi + proj.numel() * proj.element_size() or hi >= lo + ( + index_q.shape[0] * index_q.stride(0) * index_q.element_size() + ), "index_q aliases proj" + assert position_ids.shape == (3, num_tokens) and position_ids.is_contiguous(), ( + f"position_ids must be a contiguous [3, {num_tokens}]; got {tuple(position_ids.shape)}" + ) + assert position_ids.dtype == torch.int32, position_ids.dtype + assert cos_sin.dtype == torch.float32 and cos_sin.shape[1] == rotary_dim, cos_sin.shape + assert rotary_dim <= min(head_dim, index_head_dim) and rotary_dim % 2 == 0, rotary_dim + if num_tokens: + assert int(position_ids.min()) >= 0, "negative position" + assert int(position_ids.max()) < cos_sin.shape[0], "position past the rope table" + + +qwen4_exp_qk_norm_rope = Qwen4ExpQkNormRope() diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_scores.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_scores.py index c30b7da0d376..89e2c0ad7711 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_scores.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_scores.py @@ -51,7 +51,7 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: #: The callable's own threshold: at or below it every call is row-wise. DECODE_ROW_LIMIT = 256 #: Index heads must divide this into more than one row for the query tile. -DOT_TILE_WIDTH = 16 +DOT_TILE_WIDTH = 128 CAPTURE_PAGES = 260 @@ -360,6 +360,78 @@ def test_both_modes_agree_on_the_visible_columns() -> None: assert torch.equal(padded[visible], unspecified[visible]) +def test_pdl_launch_is_bitwise_the_plain_one_and_reads_after_the_wait() -> None: + """`use_pdl` changes the launch, not the result; under graph replay a producer rewrites + q right before the op, and a read of it before the wait would see the previous + replay's.""" + args = _build(_case(4, 4, only_visible_blocks=True), seed=313) + visible = _visible(args) + plain = op(**args) + assert torch.equal(op(**{**args, "use_pdl": True})[visible], plain[visible]) + src = args["q"].clone() + + def step(): + args["q"].copy_(src * 1.0) + return op(**{**args, "use_pdl": True}) + + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + step() + torch.cuda.current_stream().wait_stream(warmup) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_out = step() + for trial in range(3): + src.copy_(_build(_case(4, 4, only_visible_blocks=True), seed=314 + trial)["q"]) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(graph_out[visible], op(**args)[visible]), f"replay {trial}" + + +def test_prefetch_reads_the_fresh_anchor_key_and_q_after_the_wait() -> None: + """`prefetch` reads pages and keys before the PDL wait. Under graph replay a producer + rewrites q and, at rows whose position ends a group, the last visible block's anchor + key -- what the decode prologue writes -- right before the op; both must be seen. + The other keys are static, as they are in the decode step.""" + args = _build(_case(4, 4, only_visible_blocks=True), seed=411) + cr, tpb = args["compress_ratio"], args["tokens_per_block"] + pos = args["query_positions"].clone() + assert bool((pos >= cr).all()) + ends = torch.tensor([True, False, True, True], device=DEV) + args["query_positions"] = torch.where(ends, pos - (pos + 1) % cr, pos).to(pos.dtype) + pos = args["query_positions"].long() + assert bool(((pos + 1) % cr == 0).eq(ends).all()) + pages = args["block_table"][args["request_indices"].long(), pos // tpb].long() + anchors = (pages[ends], (pos % tpb)[ends]) + visible = _visible(args) + plain = op(**args) + assert torch.equal(op(**{**args, "use_pdl": True, "prefetch": True})[visible], plain[visible]) + src_q = args["q"].clone() + src_k = args["index_cache"][anchors].clone() + + def step(): + args["q"].copy_(src_q * 1.0) + args["index_cache"][anchors] = src_k * 1.0 + return op(**{**args, "use_pdl": True, "prefetch": True}) + + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + step() + torch.cuda.current_stream().wait_stream(warmup) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_out = step() + for trial in range(3): + fresh = _build(_case(4, 4, only_visible_blocks=True), seed=412 + trial) + src_q.copy_(fresh["q"]) + src_k.copy_(torch.randn(src_k.shape, device=DEV).to(src_k.dtype)) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(graph_out[visible], op(**args)[visible]), f"replay {trial}" + + def test_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 diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_store.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_store.py index 32b4c1aae597..168be9db6969 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_store.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_store.py @@ -7,9 +7,7 @@ each live row's pages drawn from a seeded permutation of the pool, `-1` past its allocation, rows past the batch holding a previous step's plausible page ids. Both planes are pre-filled with a sentinel so a stray write shows. - -Out-of-range pages, requests and slots are deliberately absent: each is a CUDA -device-side assert that poisons the process for every test after it. +A `hole` cell appends one row that resolves to no slot. """ import zlib @@ -45,7 +43,7 @@ def _compare(out, ref) -> None: DEV = "cuda" KEY_FILL = -7.0 POS_FILL = -7 -ENGINE_MAX_NUM_TOKENS = 8192 +ENGINE_MAX_NUM_TOKENS = 16384 def _bits(tensor: torch.Tensor) -> torch.Tensor: @@ -53,6 +51,27 @@ def _bits(tensor: torch.Tensor) -> torch.Tensor: return tensor.contiguous().view(as_int) +def _hole(spec, table, live_requests, allocated_cols): + """One row that resolves to no slot, and the table edit it needs.""" + tpb, width, slots, pages = (spec[k] for k in ("tokens_per_block", "width", "slots", "pages")) + kind = spec["hole"] + if kind == "unallocated_page": + assert int(table[0, allocated_cols[0]]) == -1 + return 0, allocated_cols[0] * tpb + if kind == "page_past_pool": + table[live_requests] = -1 + table[live_requests, 0] = pages + 7 + return live_requests, 0 + if kind == "request_past_table": + return slots, 0 + if kind == "negative_request": + return -1, 0 + if kind == "negative_position": + return 0, -1 + assert kind == "column_past_table" + return 0, width * tpb + + def _shared_page_planes(spec, extras: dict): """The NVFP4-KV pool's side planes: each physical page holds the `layers` index-K pages and then the position page, so the index cache (the last @@ -87,24 +106,38 @@ def _build(spec, seed: int): table = torch.randint(0, pages, (slots, width), generator=gen, dtype=torch.int32) order = torch.randperm(pages, generator=gen).to(torch.int32) - used = 0 + used, allocated_cols = 0, [] for r, (start, n) in enumerate(segments): cols = (start + n - 1) // tpb + 1 table[r] = -1 table[r, :cols] = order[used : used + cols] used += cols + allocated_cols.append(cols) assert used <= pages and len(segments) < slots, "pool or table too small for the batch" + if spec["hole"] is not None: + request, position = _hole(spec, table, len(segments), allocated_cols) + requests.append(request) + positions.append(position) + rows += 1 table = table.to(DEV) if spec["table_layout"] == "column_slice": wide = torch.full((slots, width + 4), -1, dtype=torch.int32, device=DEV) wide[:, 2 : 2 + width] = table table = wide[:, 2 : 2 + width] assert not table.is_contiguous() + elif spec["table_layout"] == "transposed": + table = table.t().contiguous().t() + assert table.stride(1) != 1 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"] == "fused_row": + # The target's fused input projection row, [q | k | v | gate | index q, k | + # pad] = 110 heads of 128: the raw index key is head 108. + row = torch.randn(rows, 110, head_dim, generator=gen, dtype=torch.float32) + token_k = row.to(device=DEV, dtype=spec["key_dtype"])[:, 108 : 108 + hkv] if spec["key_layout"] == "contiguous": token_k = token_k.contiguous() @@ -213,14 +246,6 @@ def test_only_the_named_slots_change() -> None: 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) @@ -230,30 +255,6 @@ def test_coordinates_wider_than_int32_wrap_rather_than_raise() -> None: 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) @@ -331,94 +332,87 @@ def test_the_stored_key_is_raw_not_normalized() -> None: assert distance > 0.25, f"a normed key differs from the raw one by only {distance}" -# ---------------------------------------------------------------- the guards +# ------------------------------------------------------- loud rejections # -# Each case is shown first unguarded -- accepted, and silently wrong -- then -# refused under `validating`. +# Each is a shape the kernel would otherwise read past or write wrongly; the +# entry refuses it before launching, so `is_valid` has nothing left to guard. -def _unguarded_then_refused(args, match: str) -> None: - op(**args) - with pytest.raises(AssertionError, match=match): - with validating(op): - op(**args) +def _rejected(args, match: str) -> None: + before = (args["index_cache"].clone(), args["position_cache"].clone()) + with pytest.raises(ValueError, match=match): + op(**args) + assert torch.equal(_bits(args["index_cache"]), _bits(before[0])) + assert torch.equal(args["position_cache"], before[1]) -def test_guard_refuses_two_dimensional_request_indices() -> None: +def test_rejects_two_dimensional_request_indices() -> None: args = _small(seed=60) args["request_indices"] = args["request_indices"][:, None] - _unguarded_then_refused(args, "one-dimensional") + _rejected(args, "one-dimensional") -def test_guard_refuses_metadata_of_unequal_length() -> None: - """A length-1 position vector collapses every row onto one slot.""" +def test_rejects_metadata_of_unequal_length() -> None: 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) + _rejected(args, "equally long") -def test_guard_refuses_a_single_token_k_row() -> None: - """One key row broadcasts into every destination slot.""" +def test_rejects_float_metadata() -> None: 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) + args["logical_positions"] = args["logical_positions"].float() + _rejected(args, "int32 or int64") -def test_guard_refuses_a_single_coordinate_row() -> None: +def test_rejects_a_single_token_k_row() -> None: args = _small(seed=63) - args["position_coordinates"] = args["position_coordinates"][:1] - _unguarded_then_refused(args, "position_coordinates has") + args["token_k"] = args["token_k"][:1] + _rejected(args, "token_k") -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.""" +def test_rejects_a_single_coordinate_row() -> None: 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) + args["position_coordinates"] = args["position_coordinates"][:1] + _rejected(args, "position_coordinates") -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.""" +def test_rejects_a_page_size_that_is_not_the_cache_s() -> None: args = _small(seed=65) + args["tokens_per_block"] = 16 + _rejected(args, "index_cache") + + +def test_rejects_a_position_plane_with_a_different_page_size() -> None: + args = _small(seed=66) 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) + _rejected(args, "position_cache") -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.""" +def test_rejects_a_position_plane_narrower_than_the_coordinates() -> None: args = _small(seed=67) + args["position_cache"] = torch.full((8, 32, 2), POS_FILL, dtype=torch.int32, device=DEV) + _rejected(args, "position_cache") + + +def test_rejects_a_cache_with_more_heads_than_the_key() -> None: + args = _small(seed=68) args["index_cache"] = torch.full((8, 32, 2, 128), KEY_FILL, device=DEV).to(torch.bfloat16) + _rejected(args, "index_cache") + + +def test_rejects_a_one_dimensional_table() -> None: + args = _small(seed=69) + args["block_table"] = args["block_table"][0] + _rejected(args, "two-dimensional") + + +def test_huge_int64_position_is_skipped_not_truncated() -> None: + """`2^32 + 5` truncated to int32 would be position 5 of column 0; the + kernel keeps int64 and finds it past the table.""" + args = _small(seed=70) + args["logical_positions"][0] = 2**32 + 5 + before = (args["index_cache"].clone(), args["position_cache"].clone()) 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) + assert torch.equal(_bits(args["index_cache"][page, 5]), _bits(before[0][page, 5])) + assert torch.equal(args["position_cache"][page, 5], before[1][page, 5]) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store.py index a77ccf3190ef..5fed6e17c22a 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store.py @@ -97,10 +97,14 @@ def _build(spec, seed: int): rows = len(positions) qkv = torch.randn( - rows, spec["heads_q"] + 2 * hkv, head_dim, generator=gen, dtype=torch.float32 + rows, + spec["heads_q"] + 2 * hkv + spec.get("row_pad_heads", 0), + head_dim, + generator=gen, + dtype=torch.float32, ).to(device=DEV, dtype=spec["row_dtype"]) k = qkv[:, spec["heads_q"] : spec["heads_q"] + hkv] - v = qkv[:, spec["heads_q"] + hkv :] + v = qkv[:, spec["heads_q"] + hkv : spec["heads_q"] + 2 * hkv] if spec["values"] == "every_bf16": patterns = torch.arange(-(2**15), 2**15, dtype=torch.int32).to(torch.int16) patterns = patterns.view(torch.bfloat16).reshape(rows, hkv, head_dim).to(DEV) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa.py index 4f94114d1d9c..ad2dd5c7d32a 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa.py @@ -20,6 +20,9 @@ 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, +) 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, ) @@ -179,8 +182,9 @@ def _build(spec, seed): 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) + if spec["q_layout"] in ("qkv_slice", "fused_row"): + row_heads = q_heads + 2 * kv if spec["q_layout"] == "qkv_slice" else 55 + qkv = torch.randn(rows, row_heads, d, generator=gen, device=DEV).to(q_dtype) q = qkv[:, :q_heads] else: q = torch.randn(rows, q_heads, d, generator=gen, device=DEV).to(q_dtype) @@ -195,6 +199,15 @@ def _build(spec, seed): request_indices[2] = slots + 67 # a request the table does not have block_table[:, 0] = -1 # every request's first page unmapped + gate = None + if spec.get("gate") == "fused_row": + # The gate columns of the target's fused projection row (55 heads of + # D): heads 28..51, after q, k and v. + fused = torch.randn(rows, 55 * d, generator=gen, device=DEV).to(q_dtype) + gate = fused[:, 28 * d : (28 + q_heads) * d] + elif spec.get("gate") == "heads": + gate = torch.randn(rows, q_heads, d, generator=gen, device=DEV).to(q_dtype) + args = dict( q=q, k_cache=k_cache, @@ -207,6 +220,8 @@ def _build(spec, seed): ) if spec["prefill_pair"]: args.update(query_positions=positions, compress_ratio=spec["compress_ratio"]) + if gate is not None: + args.update(output_gate=gate) return args, generation @@ -369,6 +384,37 @@ def test_calls_are_bitwise_deterministic() -> None: assert torch.equal(op(**args), op(**args)) +@pytest.mark.parametrize("rows", [1, 4, 128]) +def test_output_gate_in_the_merge_is_the_standalone_gate(rows) -> None: + """The gate carried by the merge equals the ungated call followed by the + standalone gate kernel, bit for bit, and the gate moves the output.""" + spec = _with(segments=((4000, 1),) * rows, slots=rows + 32, q_layout="fused_row") + spec["gate"] = "fused_row" + args, _ = _build(spec, seed=8900 + rows) + gate = args.pop("output_gate") + assert op.fuses_output_gate(args["q"], args["k_cache"], args["selected_tokens"], gate) + ungated = op(**args) + gated = op(**args, output_gate=gate) + standalone = fused_sigmoid_mul_inplace(ungated.reshape(rows, -1).clone(), gate) + assert torch.equal(gated.reshape(rows, -1), standalone) + assert (gated.float() - ungated.float()).abs().max().item() > 0.1 + + +def test_output_gate_needs_a_merge() -> None: + """One split (300 rows) and the prefill kernel have no merge to carry the + gate: `fuses_output_gate` says so and the call refuses it.""" + for spec in ( + _with(segments=((600, 1),) * 300, slots=384), + _with(segments=((2032, 272),), prefill_pair=True), + ): + spec["gate"] = "heads" + args, _ = _build(spec, seed=9200) + gate = args["output_gate"] + assert not op.fuses_output_gate(args["q"], args["k_cache"], args["selected_tokens"], gate) + with pytest.raises(ValueError, match="output-gate fusion requires a compatible split-K"): + 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"): diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_decode.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_decode.py new file mode 100644 index 000000000000..bf9206712e68 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_decode.py @@ -0,0 +1,163 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qsa_paged_sparse_gqa_decode entry: cells against the composition +entry's fp32 reference in that entry's band (4 ulp of each head row's magnitude, 2 ulp +relative RMS), agreement with the composition entry (`qsa_paged_sparse_gqa`, which this +replaces at decode) in the same band, stale-read checks under graph replay with PDL, 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.attention.qsa_paged_sparse_gqa import ( # noqa: E501 + qsa_paged_sparse_gqa, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_sparse_gqa_decode import ( # noqa: E501 + qsa_paged_sparse_gqa_decode as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_topk_expand_decode import ( # noqa: E501 + qsa_topk_expand_decode, +) + +__extra_import_path__ = ["../../../.."] +from _validating import ( # noqa: E402,F401 — path above; fixture by name + assert_within_ulp, + certifying, + validating, +) + +assert torch.cuda.is_available(), "qsa_paged_sparse_gqa_decode requires a CUDA device" + +DEV = "cuda" +Q_HEADS, KV_HEADS, D, TPB, CR, TOKEN_TOPK, PAGES, SLOTS = 24, 2, 256, 32, 4, 2048, 320, 6 + + +def _build(spec, seed): + """Paged K/V for one decode row per position, the selection QSA's expansion emits + (512 random complete groups, then the incomplete group, then -1), q and the gate as + column views of the target's 55-head fused projection row.""" + g = torch.Generator(device=DEV).manual_seed(seed) + pos = torch.tensor(spec["pos"], dtype=torch.int64, device=DEV) + rows = pos.numel() + pages_per = [int(p) // TPB + 1 for p in spec["pos"]] + pool_pages = sum(pages_per) + 5 + page_ids = torch.randperm(pool_pages, generator=g, device=DEV).to(torch.int32) + block_table = torch.full((SLOTS, PAGES), -1, dtype=torch.int32, device=DEV) + slot_of = list(range(1, rows + 1)) + cursor = 0 + for slot, count in zip(slot_of, pages_per): + block_table[slot, :count] = page_ids[cursor : cursor + count] + cursor += count + pool = torch.randn(pool_pages, 2, KV_HEADS, TPB, D, generator=g, device=DEV).to(torch.bfloat16) + k_cache, v_cache = pool[:, 0], pool[:, 1] + fused = torch.randn(rows, 55 * D, generator=g, device=DEV).to(torch.bfloat16) + q = fused.view(rows, 55, D)[:, :Q_HEADS] + gate = fused[:, 28 * D : (28 + Q_HEADS) * D] if spec["gated"] else None + block_topk = TOKEN_TOPK // CR + selected = torch.full((rows, TOKEN_TOPK + CR - 1), -1, dtype=torch.int32, device=DEV) + for r, p in enumerate(spec["pos"]): + visible = p + 1 + groups = visible // CR + chosen = torch.randperm(groups, generator=g, device=DEV)[: min(groups, block_topk)] + toks = (chosen[:, None] * CR + torch.arange(CR, device=DEV)).reshape(-1) + tail = torch.arange(groups * CR, visible, device=DEV) + row = torch.cat([toks, tail]).to(torch.int32) + selected[r, : row.numel()] = row + if spec.get("holes"): + selected[r, 1 : row.numel() : 2] = -1 + request_indices = torch.tensor(slot_of, dtype=torch.int32, device=DEV) + return [q, k_cache, v_cache, block_table, selected, request_indices, TPB, D**-0.5, gate] + + +def _kw(args): + names = ( + "q", + "k_cache", + "v_cache", + "block_table", + "selected_tokens", + "request_indices", + "tokens_per_block", + "softmax_scale", + "output_gate", + ) + return dict(zip(names, args)) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) + out = op(*args) + ref = op.reference(*args) + assert out.dtype == torch.bfloat16 and out.shape == ref.shape and out.is_contiguous() + assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0) + + +@pytest.mark.parametrize("gated", [True, False]) +def test_matches_the_composition(gated) -> None: + args = _build(dict(pos=(8700,), gated=gated), seed=91) + got = op(*args) + want = qsa_paged_sparse_gqa(**_kw(args)) + assert_within_ulp(got, want, element_ulp=4.0, rms_ulp=2.0) + print(f"\n gated={gated}: bitwise the composition's {(got == want).float().mean().item():.4f}") + + +def _replays_like_eager(step, refresh) -> None: + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + step() + torch.cuda.current_stream().wait_stream(warmup) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_out = step() + for trial in range(3): + refresh() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(graph_out, step()), f"replay {trial}" + + +def test_reads_q_and_the_gate_fresh_under_graph_replay() -> None: + """A plain launch rewrites q and the gate (one fused row) inside the graph right before + the op. The op reads both before its waits, which is sound only because such a producer + releases the op at its end: this pins that contract.""" + args = _build(dict(pos=(8700,), gated=True), seed=17) + fused = args[0]._base if args[0]._base is not None else args[0] + src = fused.clone() + + def step(): + fused.copy_(src * 1.0) + return op(*args) + + _replays_like_eager(step, lambda: src.copy_(torch.randn(src.shape, device=DEV).to(src.dtype))) + + +def test_reads_the_selection_after_the_wait_under_graph_replay() -> None: + """The target's top-k runs inside the graph right before the op: a PDL producer that + releases the op right after its own wait, before it writes the selection. A read of the + selection before the op's wait would see the previous replay's.""" + args = _build(dict(pos=(8700,), gated=True), seed=23) + pos = torch.tensor([8700], dtype=torch.int64, device=DEV) + ends = ((pos + 1) // CR).to(torch.int32) + seqlen = (pos + 1).to(torch.int32) + scores = torch.rand(1, 2304, device=DEV) + topk = torch.empty(1, TOKEN_TOPK // CR, dtype=torch.int32, device=DEV) + + def step(): + sel = qsa_topk_expand_decode(scores, ends, pos, seqlen, topk, CR, TOKEN_TOPK) + return op(*args[:4], sel, *args[5:]) + + _replays_like_eager(step, lambda: scores.copy_(torch.rand(scores.shape, device=DEV))) + + +def test_guards() -> None: + args = _build(dict(pos=(8700,), gated=True), seed=5) + many = [a.expand(300, *a.shape[1:]) if i in (0, 4, 5) else a for i, a in enumerate(args[:6])] + with validating(op), pytest.raises(AssertionError): + op(*many, *args[6:8], None) # 300 rows: past MAX_ROWS + with validating(op), pytest.raises(AssertionError): + op(*args[:8], args[8].float()) # gate not in q's dtype diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py new file mode 100644 index 000000000000..3d6a89a2c20b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py @@ -0,0 +1,468 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qsa_sparse_gqa_rowtile entry. + +Same paged state as `qsa_paged_sparse_gqa`'s test: one `[pages, 2, kv, tpb, D]` +HND pool, the engine's slot table (more slots than requests, a seeded +permutation, `-1` past each sequence), and selection rows laid out the way +QSA's top-k + expansion emit them (complete groups expanded to tokens, the +trailing partial group, then `-1`, packed at the front). The band and the +reference are that entry's: this op replaces it on context launches. +""" + +import math +import re + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention import ( # noqa: E501 + qsa_sparse_gqa_rowtile as _module, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_sparse_gqa import ( # noqa: E501 + qsa_paged_sparse_gqa as per_row, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_sparse_gqa_rowtile import ( # noqa: E501 + qsa_sparse_gqa_rowtile as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401 + +assert torch.cuda.is_available(), "qsa_sparse_gqa_rowtile requires a CUDA device" + +DEV = "cuda" + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """`qsa_paged_sparse_gqa`'s band: 4 ulp of each (row, head)'s magnitude, + 2 ulp relative RMS, in q's dtype, against the fp32 reference rounded once.""" + assert out.dtype == ref.dtype, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + if out.numel() == 0: + return + assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0) + + +class _LaunchRecorder: + """Records which of the entry's kernels a call indexed with a grid; the bitmap + pass (`_row_bits`, a plain call) is recorded with its one-CTA-a-row grid.""" + + def __init__(self, inner, name, log): + self._inner, self._name, self._log = inner, name, log + + def __getitem__(self, grid): + self._log.append((self._name, tuple(grid))) + return self._inner[grid] + + def __call__(self, selected_tokens, *args, **kwargs): + self._log.append((self._name, (selected_tokens.shape[0],))) + return self._inner(selected_tokens, *args, **kwargs) + + +_KERNELS = ( + "_row_bits", + "_rowtile_plan_kernel", + "_rowtile_lone_kernel", + "_rowtile_attention_kernel", +) +#: The row-tile pass on SM103 tensor cores, which stands in for `_rowtile_attention_kernel` +#: on eligible launches (bf16 q and K/V, D = 256). +_CUTE_PASS = "_rowtile_attention_cute" + + +def _expected_kernels(args) -> list: + cute = _module._cute_eligible(args["q"], args["k_cache"], args["v_cache"], 128, 64) + return [*_KERNELS[:3], _CUTE_PASS if cute else _KERNELS[3]] + + +@pytest.fixture +def launches(monkeypatch): + log: list = [] + for name in (*_KERNELS, _CUTE_PASS): + inner = getattr(_module, name) + monkeypatch.setattr(_module, name, _LaunchRecorder(inner, name, log)) + return log + + +def _selected_tokens(positions, seed, compress_ratio, token_topk, shared_scores=False): + """QSA's expansion layout: a random `min(available, block_topk)`-subset of + the complete visible groups, expanded, then the trailing partial group, + then `-1`, packed at the front. With `shared_scores` every row ranks the + groups by one score vector, so adjacent rows pick nearly the same set (a + sparse union, as on real selections).""" + block_topk = token_topk // compress_ratio + final_topk = token_topk + compress_ratio - 1 + gen = torch.Generator(device=DEV).manual_seed(seed) + rows = positions.numel() + visible = positions.long() + 1 + available = visible // compress_ratio + taken = available.clamp(max=block_topk) + width = max(block_topk, int(available.max().item()) if rows else 0) + out = torch.empty(rows, final_topk, dtype=torch.int32, device=DEV) + column = torch.arange(final_topk, device=DEV)[None, :] + for start in range(0, rows, 2048): + stop = min(start + 2048, rows) + avail, take, vis = available[start:stop], taken[start:stop], visible[start:stop] + if shared_scores: + scores = torch.rand(1, width, generator=gen, device=DEV).expand(stop - start, -1) + gen.manual_seed(seed) + else: + scores = torch.rand(stop - start, width, generator=gen, device=DEV) + group_ids = torch.arange(width, device=DEV)[None, :] + scores = torch.where(group_ids < avail[:, None], scores, torch.inf) + chosen = scores.argsort(dim=1)[:, :block_topk] + group_of_column = (column // compress_ratio).clamp(max=block_topk - 1) + body = torch.gather(chosen, 1, group_of_column.expand(stop - start, -1)) + body = body * compress_ratio + column % compress_ratio + in_body = column < (take * compress_ratio)[:, None] + tail_column = column - (take * compress_ratio)[:, None] + tail = (avail * compress_ratio)[:, None] + tail_column + in_tail = ~in_body & (tail_column >= 0) & (tail < vis[:, None]) + out[start:stop] = torch.where(in_body, body, torch.where(in_tail, tail, -1)).to(torch.int32) + return out + + +def _build(spec, seed): + """Real paged state for one spec, keyed by `__call__`'s parameter names.""" + gen = torch.Generator(device=DEV).manual_seed(seed) + tpb, kv, d, q_heads = ( + spec["tokens_per_block"], + spec["kv_heads"], + spec["head_dim"], + spec["q_heads"], + ) + segments = list(spec["segments"]) + positions = torch.cat([torch.arange(s, s + n, device=DEV) for s, n in segments]) + req = torch.cat([torch.full((n,), i, device=DEV) for i, (_, n) in enumerate(segments)]) + rows, slots = positions.numel(), spec["slots"] + assert len(segments) <= slots + slot_of = torch.randperm(slots, generator=torch.Generator().manual_seed(seed))[: len(segments)] + 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"], + spec.get("shared_scores", False)) # fmt: skip + 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 + return dict( + q=q, + k_cache=k_cache, + v_cache=v_cache, + block_table=block_table, + selected_tokens=selected, + request_indices=request_indices, + tokens_per_block=tpb, + softmax_scale=d**-0.5, + query_positions=positions.to(spec["request_dtype"]), + compress_ratio=spec["compress_ratio"], + ) + + +def _bound(positions, spec): + cr = spec["compress_ratio"] + visible = positions.long() + 1 + return (visible // cr).clamp(max=spec["token_topk"] // cr) * cr + visible % cr + + +def _is_prefix_packed(table): + valid = table >= 0 + columns = torch.arange(table.shape[1], device=table.device)[None, :] + return (valid == (columns < valid.sum(1)[:, None])).all(1) + + +def _shared_groups(request_indices, rows_per_tile=10): + """How many 10-row groups name a single request (served by the row-tile pass).""" + rows = request_indices.numel() + padded = torch.full((math.ceil(rows / rows_per_tile) * rows_per_tile,), -1, device=DEV) + padded[:rows] = request_indices.long() + grid = padded.view(-1, rows_per_tile) + live = torch.arange(padded.numel(), device=DEV).view(-1, rows_per_tile) < rows + low = torch.where(live, grid, 1 << 40).min(1).values + high = torch.where(live, grid, -1).max(1).values + return int((low == high).sum()) + + +@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 not spec.get("holes"): + # The per-row kernel's precondition, which this op shares. + selected = args["selected_tokens"] + assert bool(_is_prefix_packed(selected).all()) + assert torch.equal((selected >= 0).sum(1), _bound(args["query_positions"], spec)) + out = op(**args) + assert [name for name, _ in launches] == _expected_kernels(args) + groups = math.ceil(rows / 10) + assert launches[0][1] == (rows,) and launches[1][1] == (groups,) + assert launches[2][1] == (min(rows, 256), spec["kv_heads"]) + assert launches[3][1] == (groups, spec["kv_heads"]) + assert out.shape == (rows, spec["q_heads"], spec["head_dim"]) + assert out.dtype == args["q"].dtype and out.is_contiguous() + assert not bool(out.isnan().any()) + _compare(out, op.reference(**args)) + + +# -------------------------------------------------------------------------- +# semantics a cell cannot state +# -------------------------------------------------------------------------- + +CKPT = op.CELLS[0].spec +assert (CKPT["q_heads"], CKPT["kv_heads"], CKPT["head_dim"]) == (24, 2, 256) + + +def _with(segments, **overrides): + return {**CKPT, "segments": tuple(segments), **overrides} + + +def _plant_past_the_bound(args): + """Put each row's own newest token in its first unscanned column (bound + rounded up to 64), where the row has one.""" + positions = args["query_positions"] + final_topk = args["selected_tokens"].shape[1] + bound = _bound(positions, CKPT).clamp(max=final_topk) + first = (bound + 63) // 64 * 64 + gap_rows = (first < final_topk).nonzero().flatten() + assert gap_rows.numel() > 0 + planted = args["selected_tokens"].clone() + planted[gap_rows, first[gap_rows]] = positions[gap_rows].to(torch.int32) + return planted, gap_rows + + +@pytest.mark.parametrize( + "segments", + [ + [(0, 8191)], + [(0, 900), (907, 200), (1814, 150), (2721, 28)] + + [(100 + (499 * g) % 8000, 1) for g in range(28)], + ], + ids=["shared-groups", "mixed-with-lone-rows"], +) +def test_prefill_bound_is_honoured(segments, launches) -> None: + """An index past the per-row kernel's bound is not attended, on the + row-tile pass and on the lone-row pass alike; the planted index moves the + right answer by more than 0.1, so the comparison is not vacuous.""" + args = _build(_with(segments, slots=64), seed=11000) + planted, gap_rows = _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[gap_rows].float() - without_it[gap_rows].float()).abs().max().item() > 0.1 + assert _shared_groups(args["request_indices"]) > 0 + + +def test_repeated_index_on_a_shared_group_is_attended_once() -> None: + """The row-tile pass attends to the set a row names (the documented + precondition): a repeated column matches the reference without the + repeat. Rows of a group that spans requests (rows 10..15 here: 15 rows of + one request, then one of another) keep per-column semantics.""" + args = _build(_with([(900, 15), (5000, 1)], slots=8), seed=8700) + selected = args["selected_tokens"].clone() + once = selected.clone() + shared_row, lone_row = 3, 15 + for r in (shared_row, lone_row): + assert int(selected[r, 0]) >= 0 + last = int((selected[r] >= 0).sum()) - 1 + selected[r, last] = selected[r, 0] # the row's last token becomes a repeat of its first + once[r, last] = -1 + out = op(**{**args, "selected_tokens": selected}) + _compare(out[:10], op.reference(**{**args, "selected_tokens": once})[:10]) + _compare(out[10:], op.reference(**{**args, "selected_tokens": selected})[10:]) + + +def test_matches_the_per_row_entry_on_dumped_layout() -> None: + """Same arguments, same answer as `qsa_paged_sparse_gqa` within twice the band.""" + args = _build( + _with([(0, 1200), (3000, 900)] + [(8300 + g, 1) for g in range(9)], slots=16), seed=4242 + ) + out, theirs = op(**args), per_row(**args) + assert_within_ulp(out, theirs, element_ulp=8.0, rms_ulp=4.0) + + +def test_calls_are_bitwise_deterministic() -> None: + args = _build(_with([(0, 2000), (500, 300), (7000, 1)], slots=8), seed=8800) + assert torch.equal(op(**args), op(**args)) + + +def test_zero_rows_return_an_empty_output() -> None: + args = _build(_with([(900, 4)], slots=8), seed=8400) + for name in ("q", "selected_tokens", "request_indices", "query_positions"): + args[name] = args[name][:0] + out = op(**args) + assert out.shape == (0, CKPT["q_heads"], CKPT["head_dim"]) and out.dtype == torch.bfloat16 + + +_REJECTIONS = { + "no_prefill_pair": ( + lambda a: {**a, "query_positions": None, "compress_ratio": None}, + "the row-tile kernel needs the prefill pair", + ), + "float32_query": (lambda a: {**a, "q": a["q"].float()}, "QSA queries must use BF16 or FP16"), + "int64_selection": ( + lambda a: {**a, "selected_tokens": a["selected_tokens"].long()}, + "QSA selected tokens must be a two-dimensional int32 tensor", + ), + "width_vs_ratio": ( + lambda a: {**a, "selected_tokens": a["selected_tokens"][:, :64].contiguous()}, + "incompatible with its compression ratio", + ), + "strided_request_indices": ( + lambda a: {**a, "request_indices": torch.zeros(8, dtype=torch.int32, device=DEV)[::2]}, + "QSA request indices must be contiguous", + ), + "host_query": (lambda a: {**a, "q": a["q"].cpu()}, "needs every tensor on one CUDA device"), +} + + +@pytest.mark.parametrize("case", list(_REJECTIONS), ids=list(_REJECTIONS)) +def test_rejects(case) -> None: + args = _build(_with([(600, 4)], slots=8), seed=9100) + mutate, message = _REJECTIONS[case] + with pytest.raises(ValueError, match=re.escape(message)): + op(**mutate(args)) + + +class _PlanRecorder: + """Captures the plan (`group_plan`) the row-tile launch reads.""" + + def __init__(self, inner, log): + self._inner, self._log = inner, log + + def __getitem__(self, grid): + def launch(*args, **kwargs): + self._log.append(args[5].clone()) + return self._inner[grid](*args, **kwargs) + + return launch + + +def _expected_union(args, rows_per_tile=10): + """Per 10-row group: the number of distinct selected tokens within each + row's bound whose page the group's request maps, when the group names + one request (else -1), and the dense walk's tiles.""" + selected, positions = args["selected_tokens"], args["query_positions"] + request_indices, table = args["request_indices"].long(), args["block_table"] + tpb, rows = args["tokens_per_block"], selected.shape[0] + bound = _bound(positions, CKPT).clamp(max=selected.shape[1]) + columns = torch.arange(selected.shape[1], device=DEV)[None, :] + counts, dense = [], [] + for start in range(0, rows, rows_per_tile): + stop = min(start + rows_per_tile, rows) + group_requests = request_indices[start:stop].unique() + if group_requests.numel() != 1: + counts.append(-1) + dense.append(0) + continue + tokens = selected[start:stop].long() + tokens = tokens[(columns < bound[start:stop, None]) & (tokens >= 0)].unique() + pages = table[int(group_requests)][(tokens // tpb).clamp(max=table.shape[1] - 1)] + kept = tokens[(tokens // tpb < table.shape[1]) & (pages >= 0)] + counts.append(kept.numel()) + dense.append((int(tokens.max()) + 64) // 64 if tokens.numel() else 0) + return torch.tensor(counts), torch.tensor(dense) + + +@pytest.mark.parametrize( + "segments, shared_scores", + [([(4096, 4096)], True), ([(900, 40), (950, 40), (3000, 25)], False)], + ids=["sparse-union", "two-requests-unmapped-page"], +) +def test_packed_walk_is_the_resident_union(segments, shared_scores, monkeypatch) -> None: + """The row-tile pass walks exactly the union of its rows' resident + selected tokens, 64 per tile: the plan's list length is that union's + size for every one-request group, and on rows ranking one score vector + the walk is far shorter than every tile up to the group's highest token.""" + args = _build(_with(segments, slots=16, shared_scores=shared_scores), seed=7300) + if not shared_scores: + args["block_table"][:, 3] = -1 # an unmapped page inside every request's span + log: list = [] + for name in ("_rowtile_attention_kernel", _CUTE_PASS): + monkeypatch.setattr(_module, name, _PlanRecorder(getattr(_module, name), log)) + out = op(**args) + _compare(out, op.reference(**args)) + plan = log[0].cpu() + counts, dense = _expected_union(args) + one_request = counts >= 0 + assert torch.equal(plan[one_request, 1], counts[one_request].int()) + taken = plan[:, 0] >= 0 + assert bool(taken.any()) + assert torch.equal(plan[taken, 0], (plan[taken, 1] + 63) // 64) + if shared_scores: + assert int(plan[taken, 0].sum()) * 2 < int(dense[taken].sum()) + + +def _triton_bits(args): + """The Triton bitmap pass (`_rowtile_bits_kernel`) into a zeroed bitmap.""" + selected, positions = args["selected_tokens"], args["query_positions"].contiguous() + tpb, cr, width = args["tokens_per_block"], args["compress_ratio"], args["block_table"].shape[1] + rows, final_topk = selected.shape + num_words = (width * tpb + 63) // 64 + bits = torch.zeros((rows, num_words), dtype=torch.int64, device=DEV) + info = torch.empty((rows, 2), dtype=torch.int32, device=DEV) + _module._rowtile_bits_kernel[(rows,)]( + selected, positions, bits, info, selected.stride(0), selected.stride(1), bits.stride(0), + TOPK=final_topk, BLOCK_TOPK=(final_topk - cr + 1) // cr, COMPRESS_RATIO=cr, + TOKENS_PER_BLOCK=tpb, PAGE_TABLE_WIDTH=width, BOUND_TILE=64, BLOCK_C=1024, num_warps=4, + ) # fmt: skip + return bits, info + + +@pytest.mark.parametrize( + "order, holes, positions_dtype", + [ + ("emitted", False, torch.int32), + ("shuffled", True, torch.int64), + ("emitted", True, torch.int64), + ], + ids=["emitted", "shuffled-with-holes", "holes-int64-positions"], +) +def test_bitmap_pass_is_the_triton_pass_bitwise(order, holes, positions_dtype) -> None: + """The CuTe bitmap pass writes the Triton pass's `bits` and `row_info` word for word: + on the emitted token order, on columns shuffled within each row (any order is + correct; order only changes the shared-memory flushes), and with an all-padding row + and tokens past the table.""" + args = _build(_with([(0, 900), (907, 2200), (5000, 37)], slots=16), seed=8100) + args["query_positions"] = args["query_positions"].to(positions_dtype) + selected = args["selected_tokens"] + if order == "shuffled": + perm = torch.argsort(torch.rand(selected.shape, device=DEV), dim=1) + selected = torch.gather(selected, 1, perm) + if holes: + selected = selected.clone() + selected[1] = -1 + selected[5, :8] = args["block_table"].shape[1] * args["tokens_per_block"] + 3 + args["selected_tokens"] = selected.contiguous() + want_bits, want_info = _triton_bits(args) + width = args["block_table"].shape[1] + got_bits, got_info = _module._row_bits( + args["selected_tokens"], args["query_positions"], args["compress_ratio"], + args["tokens_per_block"], width, want_bits.shape[1], + ) # fmt: skip + assert torch.equal(got_info, want_info) + assert torch.equal(got_bits, want_bits) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_expand_decode.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_expand_decode.py new file mode 100644 index 000000000000..d13a4572f629 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_expand_decode.py @@ -0,0 +1,164 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qsa_topk_expand_decode entry. + +Cells against the entry's reference (a stable descending sort, ties to the higher block, +then the expansion entry's reference): the block list and the token selection exactly. +Columns past each row's visible count carry garbage (huge values, NaN) that must never +be read. Beyond the cells: bitwise the composition it replaces (`qsa_topk_select`, then +`qsa_canonical_selection`, then `qsa_expand_block_indices`) on tie-free and tie-heavy rows; +a stale-read check under graph replay. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_canonical_selection import ( # noqa: E501 + qsa_canonical_selection, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_expand_block_indices import ( # noqa: E501 + qsa_expand_block_indices, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_topk_expand_decode import ( # noqa: E501 + qsa_topk_expand_decode as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_topk_select import ( # noqa: E501 + qsa_topk_select, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "qsa_topk_expand_decode requires a CUDA device" + +DEV = "cuda" +W, K, CR, TOKEN_TOPK = 2560, 512, 4, 2048 + + +def _build(spec, seed): + g = torch.Generator(device=DEV).manual_seed(seed) + pos = torch.tensor(spec["pos"], dtype=torch.int64, device=DEV) + rows = pos.numel() + scores = torch.randn(rows, W, generator=g, device=DEV).abs() * 3.0 + if spec.get("ties"): + scores = torch.round(scores * 2.0) / 2.0 # few distinct values: ties at the threshold + if spec.get("narrow"): + scores = 1.0 + 0.05 * torch.rand(rows, W, generator=g, device=DEV) + if spec.get("const"): + scores = torch.full((rows, W), 0.75, device=DEV) + if spec.get("tied"): + # the 400 best distinct, then `tied` equal ones straddling the k-th place + scores = torch.rand(rows, W, generator=g, device=DEV) + best = torch.randperm(W // 4, generator=g, device=DEV) + scores[:, best[:400]] = 2.0 + torch.rand(rows, 400, generator=g, device=DEV) + scores[:, best[400 : 400 + spec["tied"]]] = 1.5 + ends = ((pos + 1) // CR).to(torch.int32) + for r in range(rows): # garbage past the visible blocks + n = int(ends[r]) + scores[r, n:] = 1.0e30 + if n + 3 < W: + scores[r, n + 3] = float("nan") + seqlen = (pos + 1).to(torch.int32) + topk = torch.full((rows, K), -7, dtype=torch.int32, device=DEV) + return [scores, ends, pos, seqlen, topk, CR, TOKEN_TOPK] + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + want_topk, want_sel = op.reference(*args) + sel = op(*args) + assert sel.dtype == torch.int32 and sel.shape == want_sel.shape + assert torch.equal(args[4], want_topk), "block list" + assert torch.equal(sel, want_sel.to(torch.int32)), "token selection" + + +def _rows(kind: str, rows: int, seed: int) -> torch.Tensor: + """Score rows the canonical rewrite exists for: tie-free, relu-sparse (fewer positive + scores than the budget, so the zero class straddles it), two-valued and all-equal.""" + g = torch.Generator(device=DEV).manual_seed(seed) + if kind == "tie_free": + return torch.rand(rows, W, generator=g, device=DEV) + 0.5 + if kind == "relu_sparse": + scores = torch.relu(torch.randn(rows, W, generator=g, device=DEV) - 2.2) + scores[:, ::7] = -0.0 # negative zero ties with zero + return scores + if kind == "two_valued": + return torch.where(torch.rand(rows, W, generator=g, device=DEV) < 0.3, 2.0, 1.0) + assert kind == "all_equal" + return torch.full((rows, W), 0.75, device=DEV) + + +@pytest.mark.parametrize("kind", ["tie_free", "relu_sparse", "two_valued", "all_equal"]) +def test_bitwise_the_canonical_composition(kind) -> None: + """Windows across the budget (below it, at it, past it, the whole table); the block + list equals `qsa_topk_select` + `qsa_canonical_selection` and the token selection that + list's expansion, bitwise.""" + rows = 64 + scores = _rows(kind, rows, seed=zlib.crc32(kind.encode()) % 2**31) + g = torch.Generator(device=DEV).manual_seed(5) + pos = torch.randint(0, CR * W, (rows,), generator=g, device=DEV, dtype=torch.int64) + pos[:6] = torch.tensor([0, 4 * K - 1, 4 * K, 4 * K + 3, 8700, CR * W - 1], device=DEV) + ends = ((pos + 1) // CR).to(torch.int32) + seqlen = (pos + 1).to(torch.int32) + topk = torch.full((rows, K), -7, dtype=torch.int32, device=DEV) + sel = op(scores, ends, pos, seqlen, topk, CR, TOKEN_TOPK) + starts = torch.zeros(rows, dtype=torch.int32, device=DEV) + want = torch.full_like(topk, -7) + qsa_topk_select(scores, starts, ends, want, index_topk=K) + qsa_canonical_selection(scores, starts, ends, want) + want_sel = qsa_expand_block_indices(want, pos, seqlen, compress_ratio=CR, token_topk=TOKEN_TOPK) + assert torch.equal(topk, want), "block list" + assert torch.equal(sel, want_sel.to(torch.int32)), "token selection" + + +@pytest.mark.parametrize("kind", ["tie_free", "relu_sparse"]) +def test_a_context_chunk_into_a_row_range_of_the_selection(kind) -> None: + """A context chunk's 2048 consecutive rows (one request, positions 6144..8191), written + through `out` into rows [64, 2112) of a step-sized selection: the block list and those + rows equal the canonical composition bitwise; the rows around them are untouched.""" + rows = 2048 + scores = _rows(kind, rows, seed=31) + pos = torch.arange(6144, 6144 + rows, device=DEV, dtype=torch.int64) + ends = ((pos + 1) // CR).to(torch.int32) + seqlen = torch.full((rows,), 6144 + rows, dtype=torch.int32, device=DEV) + topk = torch.full((rows, K), -7, dtype=torch.int32, device=DEV) + whole = torch.full((rows + 128, TOKEN_TOPK + CR - 1), -9, dtype=torch.int32, device=DEV) + sel = op(scores, ends, pos, seqlen, topk, CR, TOKEN_TOPK, out=whole[64 : 64 + rows]) + assert sel.data_ptr() == whole[64].data_ptr() + starts = torch.zeros(rows, dtype=torch.int32, device=DEV) + want = torch.full_like(topk, -7) + qsa_topk_select(scores, starts, ends, want, index_topk=K) + qsa_canonical_selection(scores, starts, ends, want) + want_sel = qsa_expand_block_indices(want, pos, seqlen, compress_ratio=CR, token_topk=TOKEN_TOPK) + assert torch.equal(topk, want), "block list" + assert torch.equal(whole[64 : 64 + rows], want_sel.to(torch.int32)), "token selection" + assert (whole[:64] == -9).all() and (whole[64 + rows :] == -9).all(), "rows outside out" + + +def test_reads_after_the_wait_under_graph_replay() -> None: + args = _build(dict(pos=(8700,)), seed=21) + src = args[0].clone() + + def step(): + args[0].copy_(src * 1.0) + return op(*args) + + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + step() + torch.cuda.current_stream().wait_stream(warmup) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_sel = step() + for trial in range(3): + src.copy_(_build(dict(pos=(8700,)), seed=100 + trial)[0]) + graph.replay() + torch.cuda.synchronize() + got_topk = args[4].clone() + sel = op(src.clone(), *args[1:]) + assert torch.equal(graph_sel, sel) and torch.equal(got_topk, args[4]), f"replay {trial}" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qwen4_exp_attn_decode_prologue.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qwen4_exp_attn_decode_prologue.py new file mode 100644 index 000000000000..24b1b58b3521 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qwen4_exp_attn_decode_prologue.py @@ -0,0 +1,131 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qwen4_exp_attn_decode_prologue entry: every cell bitwise the +composition it replaces (`qwen4_exp_qk_norm_rope`, `qsa_paged_kv_store`, +`qsa_paged_index_store`, `qsa_prefill_compress`) on every buffer either writes -- the +projection, the index query, the K/V pool, the index plane, the position cache; a +stale-read check under graph replay with PDL.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qwen4_exp_attn_decode_prologue import ( # noqa: E501 + qwen4_exp_attn_decode_prologue as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "qwen4_exp_attn_decode_prologue requires a CUDA device" + +DEV = "cuda" +HQ, HK, D, IH, ID, IQ_COL, W = 24, 2, 256, 4, 128, 13312, 14080 +TPB, CR, ROT, TABLE, SLOTS, WIDTH = 32, 4, 64, 16384, 8, 300 +GEOM = dict( + tokens_per_block=TPB, compress_ratio=CR, num_heads_q=HQ, num_heads_k=HK, head_dim=D, + index_heads=IH, index_head_dim=ID, index_q_col=IQ_COL, rotary_dim=ROT, eps=1e-6, + mrope_section1=11, mrope_section2=10, +) # fmt: skip + + +def _build(spec, seed): + g = torch.Generator(device=DEV).manual_seed(seed) + pos = list(spec["pos"]) + rows = len(pos) + pages = SLOTS * WIDTH // 2 + perm = torch.randperm(pages, generator=g, device=DEV).to(torch.int32) + block_table = torch.full((SLOTS, WIDTH), -1, dtype=torch.int32, device=DEV) + used = 0 + for r, p in enumerate(pos): + n = p // TPB + 1 + block_table[r + 1, :n] = perm[used : used + n] + used += n + kv = torch.randn(pages, 2, HK, TPB, D, generator=g, device=DEV).to(torch.bfloat16) + index_cache = torch.randn(pages, TPB, 1, ID, generator=g, device=DEV).to(torch.bfloat16) + position_cache = torch.randint( + 0, TABLE, (pages, TPB, 3), generator=g, device=DEV, dtype=torch.int32 + ) + request = torch.arange(1, rows + 1, dtype=torch.int32, device=DEV) + if spec.get("unmapped") is not None: + request[spec["unmapped"]] = SLOTS + 5 + logical = torch.tensor(pos, dtype=torch.int64, device=DEV) + ids = torch.stack([logical, logical + 3, logical + 7]).to(torch.int32).contiguous() + buf = torch.empty(3, 1, rows + 5, dtype=torch.int32, device=DEV) + buf[:, 0, :rows] = ids + coords = torch.transpose(buf[:, 0, :rows], 0, 1) + angles = torch.rand(TABLE, ROT // 2, generator=g, device=DEV) * 6.28 + cos_sin = torch.cat([torch.cos(angles), torch.sin(angles)], dim=1).contiguous() + proj = (torch.randn(rows, W, generator=g, device=DEV) * 2.0).to(torch.bfloat16) + index_q = torch.zeros(rows, IH + 2, ID, dtype=torch.bfloat16, device=DEV) + weights = [ + (0.3 * torch.randn(n, generator=g, device=DEV)).to(torch.bfloat16) for n in (D, D, ID, ID) + ] + return dict( + proj=proj, index_q=index_q[:, :IH], position_ids=ids, coords=coords, cos_sin=cos_sin, + q_weight=weights[0], k_weight=weights[1], index_q_weight=weights[2], index_k_weight=weights[3], + k_cache=kv[:, 0], v_cache=kv[:, 1], index_cache=index_cache, position_cache=position_cache, + request_indices=request, logical_positions=logical, block_table=block_table, **GEOM, + ) # fmt: skip + + +_WRITTEN = ("proj", "index_q", "k_cache", "v_cache", "index_cache", "position_cache") + + +def _clone(args): + out = dict(args) + kv = args["k_cache"]._base if args["k_cache"]._base is not None else None + if kv is not None: + kv2 = kv.clone() + out["k_cache"], out["v_cache"] = kv2[:, 0], kv2[:, 1] + iq = args["index_q"]._base if args["index_q"]._base is not None else args["index_q"] + iq2 = iq.clone() + out["index_q"] = iq2[:, : args["index_q"].shape[1]] + for name in ("proj", "index_cache", "position_cache"): + out[name] = args[name].clone() + return out + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) + want = _clone(args) + op.reference(**want) + got = _clone(args) + op(**got) + torch.cuda.synchronize() + for name in _WRITTEN: + assert torch.equal(got[name], want[name]), name + + +def test_reads_after_the_wait_under_graph_replay() -> None: + """The projection is rewritten inside the graph right before the op; with PDL a read of + it before the wait would see the previous replay's row.""" + args = _build(dict(pos=(8699, 8700)), seed=23) + src = args["proj"].clone() + work = _clone(args) + + def step(): + work["proj"].copy_(src * 1.0) + op(**work) + + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + step() + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + step() + for trial in range(3): + src.copy_((torch.randn(src.shape, device=DEV) * 2.0).to(src.dtype)) + before = _clone(work) + graph.replay() + torch.cuda.synchronize() + want = _clone(before) + want["proj"].copy_(src) + op.reference(**want) + for name in _WRITTEN: + assert torch.equal(work[name], want[name]), f"replay {trial}: {name}" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qwen4_exp_qk_norm_rope.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qwen4_exp_qk_norm_rope.py new file mode 100644 index 000000000000..2cf8984e87c5 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qwen4_exp_qk_norm_rope.py @@ -0,0 +1,363 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qwen4_exp_qk_norm_rope entry.""" + +import functools +import sys +import zlib + +import pytest +import torch + +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.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qwen4_exp_qk_norm_rope import ( # noqa: E501 + CKPT, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qwen4_exp_qk_norm_rope import ( # noqa: E501 + qwen4_exp_qk_norm_rope as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import ( # noqa: E402,F401 — path above; fixture by name + assert_within_ulp, + certifying, + validating, +) + +assert torch.cuda.is_available(), "qwen4_exp_qk_norm_rope requires a CUDA device" + +#: The op's module: its two kernels' launchers and the token count between them. +M = sys.modules[type(op).__module__] + +THETA = 1e7 +EPS = 1e-6 +TABLE_ROWS = 262144 # the checkpoint's max_position_embeddings +#: Spec keys that shape the test's inputs rather than the op call. +_NOT_OP_ARGS = ("tokens", "max_position", "proj_width", "index_scratch_heads") + + +def _seed(text: str) -> int: + return zlib.crc32(text.encode()) + + +@functools.lru_cache(maxsize=1) +def _table() -> torch.Tensor: + """The target's rope table, built the way `post_load_weights` builds it.""" + rot = CKPT["rotary_dim"] + exponent = torch.arange(0, rot, 2, dtype=torch.float64, device="cuda") + inv_freq = (1.0 / (THETA ** (exponent / rot))).to(torch.float32) + positions = torch.arange(TABLE_ROWS, dtype=torch.float32, device="cuda") + angles = positions.unsqueeze(1) * inv_freq.unsqueeze(0) + return torch.cat([torch.cos(angles), torch.sin(angles)], dim=1).contiguous() + + +def _distinct_rows(tokens: int, max_position: int, g: torch.Generator) -> torch.Tensor: + """A contiguous int32 `[3, T]` whose three rows differ at every token.""" + row0 = torch.randint(0, max_position, (tokens,), generator=g, device="cuda") + d1 = torch.randint(1, max_position, (tokens,), generator=g, device="cuda") + d2 = torch.randint(1, max_position - 1, (tokens,), generator=g, device="cuda") + d2 = d2 + (d2 >= d1).long() + rows = torch.stack([row0, (row0 + d1) % max_position, (row0 + d2) % max_position]) + return rows.to(torch.int32).contiguous() + + +def _build(spec, seed, positions=None): + """`(proj, scratch, kwargs)`: the fused projection row, the index scratch + whose first `index_heads` heads are the `index_q` view, and the op's + keyword arguments. Norm weights near 1 put q/k at their largest magnitude + under the Gemma `1 + w`.""" + g = torch.Generator(device="cuda").manual_seed(seed) + t = spec["tokens"] + proj = torch.randn(t, spec["proj_width"], generator=g, device="cuda").to(torch.bfloat16) + scratch = torch.randn( + t, spec["index_scratch_heads"], spec["index_head_dim"], generator=g, device="cuda" + ).to(torch.bfloat16) + + def weight(n): + return (torch.randn(n, generator=g, device="cuda") * 0.1 + 1.0).to(torch.bfloat16) + + if positions is None: + positions = _distinct_rows(t, spec.get("max_position", 32768), g) + kwargs = {k: v for k, v in spec.items() if k not in _NOT_OP_ARGS} + kwargs.update( + index_q=scratch[:, : spec["index_heads"]], + position_ids=positions, + cos_sin=_table(), + q_weight=weight(spec["head_dim"]), + k_weight=weight(spec["head_dim"]), + index_q_weight=weight(spec["index_head_dim"]), + eps=EPS, + ) + return proj, scratch, kwargs + + +def _deviation(out: torch.Tensor, ref: torch.Tensor) -> str: + ulp = 2.0**-8 + o, r = out.float(), ref.float() + row = r.abs().amax(dim=-1, keepdim=True).clamp_min(1e-9) + elt = ((o - r).abs() / row).max().item() / ulp + rms = ((o - r).pow(2).mean().sqrt() / r.pow(2).mean().sqrt()).item() / ulp + return f"max abs {(o - r).abs().max().item():.6f}, elt {elt:.2f} ulp, rms {rms:.3f} ulp" + + +def _main_cols(spec) -> int: + return (spec["num_heads_q"] + spec["num_heads_k"]) * spec["head_dim"] + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 + certifying(op) + spec = cell.spec + proj, scratch, kwargs = _build(spec, seed=_seed(cell.why)) + before, scratch_before = proj.clone(), scratch.clone() + ref_proj, ref_iq = op.reference(proj, **kwargs) + assert torch.equal(proj, before), "the reference mutated its input" + out_proj, out_iq = op(proj, **kwargs) + assert out_proj is proj and out_iq is kwargs["index_q"] + main = _main_cols(spec) + t = spec["tokens"] + flat_ref = ref_iq.reshape(t, -1) + flat_out = out_iq.reshape(t, -1) + print(f"\n {cell.why[:50]}: main {_deviation(proj[:, :main], ref_proj[:, :main])}") + print(f" {cell.why[:50]}: index {_deviation(flat_out, flat_ref)}") + # fp32 throughout and one rounding, against an fp64 reference over the same + # table: the two can differ by one output rounding step, which is up to 2 + # ulp of the row scale (an element just above a power of two in a row whose + # largest magnitude sits just below the next one). Measured: max 1.48 ulp + # element-wise, relative RMS <= 0.011 ulp, so the RMS gate is tight. + assert_within_ulp(proj[:, :main], ref_proj[:, :main], element_ulp=2.0, rms_ulp=0.25) + assert_within_ulp(flat_out, flat_ref, element_ulp=2.0, rms_ulp=0.25) + assert torch.equal(proj[:, main:], before[:, main:]), "a column past the main k heads moved" + heads = spec["index_heads"] + assert torch.equal(scratch[:, heads:], scratch_before[:, heads:]), ( + "index_q wrote past its heads" + ) + + +def _current_pair(proj, scratch, kwargs, index_k_weight): + """What the forward ran before this entry: the CUDA norm+rope in place on + the projection's main heads, the index columns copied into the scratch, and + the CUDA norm+rope again on the scratch (row 0, plain rope).""" + spec = CKPT + tail = spec["proj_width"] // spec["head_dim"] - spec["num_heads_q"] - spec["num_heads_k"] + pos3 = kwargs["position_ids"] + fused_qk_norm_rope( + proj, num_heads_q=spec["num_heads_q"], num_heads_k=spec["num_heads_k"], + num_heads_v=tail, head_dim=spec["head_dim"], rotary_dim=spec["rotary_dim"], eps=EPS, + q_weight=kwargs["q_weight"], k_weight=kwargs["k_weight"], base=THETA, is_neox=True, + position_ids=pos3, is_qk_norm=True, use_gemma=True, use_mrope=True, + mrope_section1=spec["mrope_section1"], mrope_section2=spec["mrope_section2"], + ) # fmt: skip + flat = scratch.reshape(proj.shape[0], -1) + flat.copy_(proj[:, spec["index_q_col"] :]) + fused_qk_norm_rope( + flat, num_heads_q=spec["index_heads"], num_heads_k=1, num_heads_v=1, + head_dim=spec["index_head_dim"], rotary_dim=spec["rotary_dim"], eps=EPS, + q_weight=kwargs["index_q_weight"], k_weight=index_k_weight, base=THETA, is_neox=True, + position_ids=pos3[0].contiguous(), is_qk_norm=True, use_gemma=True, use_mrope=False, + mrope_section1=0, mrope_section2=0, + ) # fmt: skip + return proj, scratch[:, : spec["index_heads"]] + + +@pytest.mark.parametrize("tokens", [1, 19, 128]) +@pytest.mark.parametrize("rows", ["distinct", "text"]) +def test_matches_the_cuda_pair_it_replaces(tokens: int, rows: str) -> None: + """Same wiring as the two CUDA calls: mRoPE sections (11, 10) on the main + plane, coordinate 0 alone on the index plane, q/k weight assignment. Both + sides are ~1 ulp from fp64 (the CUDA op's `__sincosf` phase error grows + with position), so the band is the sum of the two entries' element bands; + a wrong wiring lands ~100x outside it.""" + spec = dict(tokens=tokens, **CKPT) + g = torch.Generator(device="cuda").manual_seed(_seed(f"pair{tokens}{rows}")) + if rows == "distinct": + positions = _distinct_rows(tokens, 9216, g) + else: + row = torch.randint(8192, 9216, (tokens,), generator=g, device="cuda", dtype=torch.int32) + positions = torch.stack([row, row, row]).contiguous() + proj, scratch, kwargs = _build(spec, seed=_seed(f"pair{tokens}{rows}"), positions=positions) + old_proj, old_scratch = proj.clone(), scratch.clone() + index_k_weight = torch.ones(CKPT["index_head_dim"], dtype=torch.bfloat16, device="cuda") + old_proj, old_iq = _current_pair(old_proj, old_scratch, kwargs, index_k_weight) + new_proj, new_iq = op(proj, **kwargs) + main = _main_cols(spec) + print(f"\n main {_deviation(new_proj[:, :main], old_proj[:, :main])}") + print(f" index {_deviation(new_iq.reshape(tokens, -1), old_iq.reshape(tokens, -1))}") + assert_within_ulp(new_proj[:, :main], old_proj[:, :main], element_ulp=3.0, rms_ulp=1.0) + assert_within_ulp( + new_iq.reshape(tokens, -1), old_iq.reshape(tokens, -1), element_ulp=3.0, rms_ulp=1.0 + ) + + +def test_index_plane_sections_select_coordinates() -> None: + """Non-zero index sections read coordinates 1 and 2: the result moves away + from the plain (0, 0) one when the rows differ.""" + spec = dict(tokens=8, **CKPT) + proj, scratch, kwargs = _build(spec, seed=_seed("index-sections")) + plain = op.reference(proj, **kwargs)[1] + kwargs.update(index_mrope_section1=11, index_mrope_section2=10) + ref = op.reference(proj, **kwargs)[1] + out = op(proj.clone(), **kwargs)[1] + assert (ref.float() - plain.float()).abs().max() > 1.0 + assert_within_ulp(out.reshape(8, -1), ref.reshape(8, -1), element_ulp=2.0, rms_ulp=0.25) + + +def test_cuda_graph_replay_reads_fresh_inputs() -> None: + """Captured once at M=4, replayed after the inputs are rewritten in place: + the replay equals an eager call on the new inputs, bit for bit.""" + spec = dict(tokens=4, **CKPT) + proj, scratch, kwargs = _build(spec, seed=_seed("graph")) + proj_new, _, kwargs_new = _build(spec, seed=_seed("graph-new")) + eager_proj, eager_scratch = proj_new.clone(), scratch.clone() + eager_kwargs = dict(kwargs_new, index_q=eager_scratch[:, : CKPT["index_heads"]]) + op(eager_proj, **eager_kwargs) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + op(proj.clone(), **dict(kwargs, index_q=scratch.clone()[:, : CKPT["index_heads"]])) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(proj, **kwargs) + proj.copy_(proj_new) + for name in ("position_ids", "q_weight", "k_weight", "index_q_weight"): + kwargs[name].copy_(kwargs_new[name]) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(proj, eager_proj) + assert torch.equal(kwargs["index_q"], eager_kwargs["index_q"]) + + +def test_zero_tokens_launches_nothing() -> None: + proj, scratch, kwargs = _build(dict(tokens=0, **CKPT), seed=1, positions=torch.empty( + 3, 0, dtype=torch.int32, device="cuda")) # fmt: skip + out_proj, out_iq = op(proj, **kwargs) + assert out_proj.shape[0] == 0 and out_iq.shape[0] == 0 + + +@pytest.mark.parametrize( + "mutate, message", + [ + (lambda kw: kw.update(index_q_col=1024), "overlap"), + (lambda kw: kw["position_ids"].__setitem__((0, 0), TABLE_ROWS), "past the rope table"), + (lambda kw: kw.update(position_ids=kw["position_ids"].t().contiguous().t()), "contiguous"), + (lambda kw: kw.update(index_q=kw["index_q"][:, :3]), ""), + ], + ids=["planes-overlap", "position-past-table", "strided-positions", "index-q-heads"], +) +def test_guard_refuses(mutate, message) -> None: + proj, _, kwargs = _build(dict(tokens=3, **CKPT), seed=_seed("guard")) + mutate(kwargs) + with validating(op), pytest.raises(AssertionError, match=message): + op(proj, **kwargs) + + +# ---- the CuTe kernel (prefill-size token counts) against the Triton kernel ---------------- + + +def _sections(kwargs): + return ( + kwargs["mrope_section1"], + kwargs["mrope_section2"], + kwargs.get("index_mrope_section1", 0), + kwargs.get("index_mrope_section2", 0), + ) + + +def _triton(proj, kwargs): + """The Triton kernel at any token count (from 65 tokens its eight-head launch).""" + M._launch_triton( + proj, kwargs["index_q"], kwargs["position_ids"], kwargs["cos_sin"], kwargs["q_weight"], + kwargs["k_weight"], kwargs["index_q_weight"], kwargs["num_heads_q"], + kwargs["num_heads_k"], kwargs["head_dim"], kwargs["index_heads"], + kwargs["index_head_dim"], kwargs["index_q_col"], kwargs["rotary_dim"], kwargs["eps"], + _sections(kwargs), + ) # fmt: skip + + +def _with_signed_zeros(proj) -> None: + """Rows whose sums and outputs turn on zero signs: +-0 and a tiny value in a rotary chunk, + -0 across a rotary/plain boundary, a whole -0 head, a zero index head.""" + proj[0, :8] = torch.tensor([0.0, -0.0, 1e-30, -0.0, 3.0, -2.0, 0.0, -0.0]) + proj[1, 300:308] = -0.0 + proj[2, 3 * 256 : 4 * 256] = -0.0 + proj[3, CKPT["index_q_col"] : CKPT["index_q_col"] + 128] = 0.0 + + +def _bitwise(a, b) -> bool: + return torch.equal(a.view(torch.int16), b.view(torch.int16)) + + +@pytest.mark.parametrize("tokens", [1024, 2048, 8193, 16384]) +@pytest.mark.parametrize("rows", ["distinct", "text"]) +def test_cute_kernel_is_the_triton_kernel_bitwise(tokens: int, rows: str) -> None: + """From `_CUTE_MIN_TOKENS` tokens the op runs the CuTe kernel: every output bit (both + planes, the untouched columns and scratch heads included) is the Triton kernel's.""" + assert tokens >= M._CUTE_MIN_TOKENS + spec = dict(tokens=tokens, **CKPT) + proj, scratch, kwargs = _build(spec, seed=_seed(f"cute{tokens}{rows}")) + if rows == "text": + kwargs["position_ids"] = kwargs["position_ids"][:1].expand(3, tokens).contiguous() + _with_signed_zeros(proj) + ref_proj, ref_scratch = proj.clone(), scratch.clone() + _triton(ref_proj, dict(kwargs, index_q=ref_scratch[:, : CKPT["index_heads"]])) + op(proj, **kwargs) + assert _bitwise(proj, ref_proj), "projection" + assert _bitwise(scratch, ref_scratch), "index scratch" + + +def test_cute_kernel_index_sections_bitwise() -> None: + """Non-zero index sections on the CuTe kernel: the index plane reads coordinates 1 and 2 + as the Triton kernel does.""" + spec = dict(tokens=2048, **CKPT) + proj, scratch, kwargs = _build(spec, seed=_seed("cute-index-sections")) + kwargs.update(index_mrope_section1=11, index_mrope_section2=10) + ref_proj, ref_scratch = proj.clone(), scratch.clone() + _triton(ref_proj, dict(kwargs, index_q=ref_scratch[:, : CKPT["index_heads"]])) + op(proj, **kwargs) + assert _bitwise(proj, ref_proj) and _bitwise(scratch, ref_scratch) + + +def test_cute_kernel_graph_replay_reads_fresh_inputs() -> None: + """The CuTe kernel captured once and replayed after the inputs are rewritten in place.""" + spec = dict(tokens=2048, **CKPT) + proj, scratch, kwargs = _build(spec, seed=_seed("cute-graph")) + proj_new, _, kwargs_new = _build(spec, seed=_seed("cute-graph-new")) + eager_proj, eager_scratch = proj_new.clone(), scratch.clone() + eager_kwargs = dict(kwargs_new, index_q=eager_scratch[:, : CKPT["index_heads"]]) + op(eager_proj, **eager_kwargs) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + op(proj.clone(), **dict(kwargs, index_q=scratch.clone()[:, : CKPT["index_heads"]])) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(proj, **kwargs) + proj.copy_(proj_new) + for name in ("position_ids", "q_weight", "k_weight", "index_q_weight"): + kwargs[name].copy_(kwargs_new[name]) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(proj, eager_proj) + assert torch.equal(kwargs["index_q"], eager_kwargs["index_q"]) + + +def test_cute_kernel_declines_strided_index_heads(monkeypatch) -> None: + """Index heads 136 elements apart are not the CuTe kernel's layout: the Triton kernel runs.""" + tokens = 2048 + spec = dict(tokens=tokens, **CKPT) + proj, _, kwargs = _build(spec, seed=_seed("cute-declines")) + wide = torch.zeros(tokens, CKPT["index_heads"], 136, dtype=torch.bfloat16, device="cuda") + kwargs["index_q"] = wide[:, :, : CKPT["index_head_dim"]] + ref_proj, ref_wide = proj.clone(), wide.clone() + _triton(ref_proj, dict(kwargs, index_q=ref_wide[:, :, : CKPT["index_head_dim"]])) + + def refuse(*args, **kwargs): + raise AssertionError("the CuTe kernel ran on strided index heads") + + monkeypatch.setattr(M, "_launch_cute", refuse) + op(proj, **kwargs) + assert _bitwise(proj, ref_proj) and _bitwise(wide, ref_wide) From ef3a73663938afec1cd5647d74c46b73742a5f7d Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Sun, 18 Oct 2026 23:00:01 -0700 Subject: [PATCH 06/28] [TRTLLM-16460][perf] Qwen3.8-Flash-Next target: Gated-DeltaNet ops and their gated norm The target's Gated-DeltaNet fast paths over an fp32 state pool, with their unit tests: - gdn_in_proj_conv1d_pdl: the decode input projection with the causal conv1d step as its epilogue, one PDL launch. - gdn_decode_gated_norm_pdl (small batches) and gdn_decode_rows_gated_norm (streamed over the state pool): one GDN decode step per row and the gated RMS norm on its output. - gdn_prefill_conv_post: the prefill causal conv1d read token-major off the input projection, with the post-conv split, Q/K L2 norm, decay gate and beta in one pass. - Updates to causal_conv1d_update / _triton_update, extract_transpose_prefill_slice, fused_gdn_post_conv, fused_sigmoid_gating_delta_rule_update and the gated norm rms_norm_gated_token_major for the calls above. Folds these commits of the review branch (their changes to the files of this commit): - qwen3.8 target: carry the decode-fusion perf work onto the feature target (62151b727b) - qwen3.8 target: GDN prefill conv pass in CuTe DSL (6925dc6dce) - qwen3.8 target: decode conv update in CuTe DSL from 64 rows (b642a1ce60) - qwen3.8 target: stream the GDN decode step with its gated norm (15dfd83ba3) - qwen3.8 target: a mixed step's generation rows on the streaming GDN kernel (73d22b699f) Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../ops/norm/rms_norm_gated_token_major.py | 66 +- .../ops/ssm/causal_conv1d_triton_update.py | 18 +- .../ops/ssm/causal_conv1d_update.py | 334 ++++- .../ssm/extract_transpose_prefill_slice.py | 19 +- .../ops/ssm/fused_gdn_post_conv.py | 16 +- .../fused_sigmoid_gating_delta_rule_update.py | 23 +- .../ops/ssm/gdn_decode_gated_norm_pdl.py | 504 +++++++ .../ops/ssm/gdn_decode_rows_gated_norm.py | 715 +++++++++ .../ops/ssm/gdn_in_proj_conv1d_pdl.py | 443 ++++++ .../ops/ssm/gdn_prefill_conv_post.py | 1285 +++++++++++++++++ ..._modeling_v2_rms_norm_gated_token_major.py | 6 + ...modeling_v2_causal_conv1d_triton_update.py | 6 + .../test_modeling_v2_causal_conv1d_update.py | 34 +- .../test_modeling_v2_fused_gdn_post_conv.py | 14 + ..._fused_sigmoid_gating_delta_rule_update.py | 15 +- ...t_modeling_v2_gdn_decode_gated_norm_pdl.py | 179 +++ ..._modeling_v2_gdn_decode_rows_gated_norm.py | 230 +++ ...test_modeling_v2_gdn_in_proj_conv1d_pdl.py | 153 ++ .../test_modeling_v2_gdn_prefill_conv_post.py | 218 +++ 19 files changed, 4252 insertions(+), 26 deletions(-) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_gated_norm_pdl.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_rows_gated_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_in_proj_conv1d_pdl.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_prefill_conv_post.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_gated_norm_pdl.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_in_proj_conv1d_pdl.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_prefill_conv_post.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/rms_norm_gated_token_major.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/rms_norm_gated_token_major.py index dca5f65af9aa..a9a667c28e0d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/rms_norm_gated_token_major.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/rms_norm_gated_token_major.py @@ -8,17 +8,49 @@ from typing import Optional import torch +import triton # Registration is lazy: the op is defined beside the mamba/GDN mixer, and # `import tensorrt_llm._torch.custom_ops` alone does not make # torch.ops.trtllm.rms_norm_gated_token_major resolve. -import tensorrt_llm._torch.modules.mamba.layernorm_gated # noqa: F401 +import tensorrt_llm._torch.modules.mamba.layernorm_gated as _layernorm_gated from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper #: The checkpoint's GDN epilogue: 48 value heads of 128, `output_gate_type: sigmoid`. _CKPT = dict(heads=48, width=128, dtype=torch.bfloat16, eps=1e-6, gate_activation="sigmoid") #: The same geometry under the op's default silu gate. _SILU = dict(_CKPT, gate_activation="silu") +#: Prefill-sized launches (at least this many rows, i.e. 683 tokens of 48 heads) run +#: the multi-row kernel at `_PREFILL_ROWS` rows a program: at 16384 tokens 94 vs 143 +#: us on GB300 (traffic roofline ~81 us), where 4 rows a program leave 196608 small +#: programs. +_PREFILL_MIN_ROWS = 32768 # capacity: fixed: a dispatch threshold, not a buffer size +_PREFILL_ROWS = 8 + + +def _prefill_multirow(x: torch.Tensor, z: torch.Tensor, gate_activation: str) -> bool: + n = x.shape[-1] + return ( + x.shape[0] >= _PREFILL_MIN_ROWS + and x.stride(-1) == 1 + and z.stride(2) == 1 + and z.stride(1) == n + and n <= 256 + and n & (n - 1) == 0 + and gate_activation in ("silu", "sigmoid") + ) + + +def _prefill_launch(x, z, weight, eps, gate_activation): + """The multi-row kernel at `_PREFILL_ROWS` rows a program (no PDL: the grid covers every SM).""" + rows, n = x.shape + out = torch.empty_like(x) + _layernorm_gated._rms_norm_gated_fwd_multirow_kernel[(triton.cdiv(rows, _PREFILL_ROWS),)]( + x, out, weight, z, None, None, x.stride(0), out.stride(0), z.stride(0), rows, eps, + N=n, ROWS=_PREFILL_ROWS, HEADS_PER_TOK=z.shape[1], GATE_SIGMOID=gate_activation == "sigmoid", + WEIGHT_IS_DELTA=False, LAUNCH_WITH_PDL=False, num_warps=4, + ) # fmt: skip + return out class RmsNormGatedTokenMajor(OpWrapper): @@ -35,6 +67,9 @@ class RmsNormGatedTokenMajor(OpWrapper): `x.dtype`, multiplied by `1 / fp8_scale[0]` and saturated at +-448 -- byte-identical to the unfused norm -> static-quantize pair. + Call site: `z` is the `[T, 6144]` column view of the GDN layer's one + `[T, 16480]` input-projection output, reshaped `[T, 48, 128]`. + Fusion boundary: norm, weight, gate and the optional quantize. The caller owns both projections producing `x` and `z`, and every reshape. """ @@ -57,6 +92,21 @@ class RmsNormGatedTokenMajor(OpWrapper): why="z a column slice: token stride 16384, 512 tokens", spec=dict(tokens=512, z="column_slice", **_CKPT), ), + *( + Cell( + why=f"z the in_proj column view: token stride 16480, {t} tokens", + spec=dict(tokens=t, z="in_proj", **_CKPT), + ) + for t in (1, 128, 8192, 16384) + ), + Cell( + why="683 tokens: the first launch at 8 rows a program", + spec=dict(tokens=683, z="in_proj", **_CKPT), + ), + Cell( + why="1001 tokens x 48 heads at 8 rows a program, silu", + spec=dict(tokens=1001, z="in_proj", **_SILU), + ), Cell( why="x row stride 2N: only stride(-1) must be 1", spec=dict(tokens=17, x="strided_rows", **_SILU), @@ -90,11 +140,13 @@ class RmsNormGatedTokenMajor(OpWrapper): 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. + program; 8 from 32768 rows, without fp8 output) when `x.stride(-1) == 1`, + `z.stride(2) == 1`, `z.stride(1) == N` and N is a power of two <= 256 -- + the checkpoint's N = 128 -- else the generic single-row kernel on a packed + copy of `z`. PDL on the 4-row path only when the grid is smaller than the + SM count and `TRTLLM_ENABLE_PDL != 0`; results are bitwise identical with + it off. The 8-row programs reduce a row's mean square in another order + than the 4-row ones (one-ulp differences on some outputs). Loud, so not guarded: 3-D `x` (ValueError), a `z` not covering `x` (AssertionError "does not match x shape"; 2-D `z` ValueError), an unknown @@ -126,6 +178,8 @@ def __call__( fp8_scale: Optional[torch.Tensor] = None, gate_activation: str = "silu", ) -> torch.Tensor: + if fp8_scale is None and z.dim() == 3 and _prefill_multirow(x, z, gate_activation): + return _prefill_launch(x, z, weight, eps, gate_activation) return torch.ops.trtllm.rms_norm_gated_token_major( x, z, weight, eps, fp8_scale, gate_activation ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_triton_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_triton_update.py index a97239b6b4b2..73c73cf96afc 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_triton_update.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_triton_update.py @@ -72,7 +72,8 @@ class CausalConv1dTritonUpdate(OpWrapper): # 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"), + # ("silu"), pad_slot_id, x_form ("token_major" | "in_proj" (token_major + # over the q|k|v columns of a [n * p, 16480] projection) | "contiguous" | "2d"), # snapshots (rows of the snapshot buffer, or 0 for none), snapshot_steps # (the buffer's step depth, default seqlen), accepted, launch_dependent_kernels. CELLS: tuple[Cell, ...] = ( @@ -116,6 +117,21 @@ class CausalConv1dTritonUpdate(OpWrapper): ), ), ) + + tuple( + Cell( + why=f"verify, x over the [n*p, 16480] in_proj, batch {batch}, {seqlen} tok/row", + spec=dict( + dim=CKPT_DIM, + batch=batch, + seqlen=seqlen, + pool_slots=ENGINE_SLOTS, + slots=_pinned(batch), + snapshots=ENGINE_SNAPSHOT_ROWS, + x_form="in_proj", + ), + ) + for batch, seqlen in ((1, 2), (32, 3)) + ) + tuple( Cell( why=f"verify, 128-slot pool, batch {batch}, {seqlen} tokens/row", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_update.py index d397c03bb956..f1d69faaaf0c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_update.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_update.py @@ -1,15 +1,27 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Single-step causal depthwise conv1d against a slot-indexed conv-state pool.""" +"""Single-step causal depthwise conv1d against a slot-indexed conv-state pool. + +The decode step's shape (one token a row, SiLU, no bias, a width-4 conv over a +3-column bf16 pool, slot indices) from `_CUTE_MIN_ROWS` rows runs a CuTe DSL +kernel (Tier 6, written for this target) that is bitwise the TRT-LLM kernel it +replaces there; every other call is `torch.ops.trtllm.causal_conv1d_update`. +""" from __future__ import annotations from typing import Optional +import cutlass +import cutlass.cute as cute import torch +from cutlass._mlir.dialects import llvm +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.cute_dsl_kernels.blackwell.utils import make_ptr #: This checkpoint's GDN `[q|k|v]` width; the conv is width 4 over a 3-column pool. CKPT_DIM = 10240 @@ -19,6 +31,283 @@ ENGINE_SLOTS = 461 +# -------------------------------------------------------------------------------------- +# The decode shape in CuTe DSL. TRT-LLM's seqlen-1 kernel runs a thread per (row, +# channel) with 2-byte accesses -- at 128 rows 1.3 M threads, four waves of one load +# latency each (7.5 us a layer on GB300). Here a thread owns 8 channels of one row: the +# weights (64 B), its state (8 channels x 3 columns, 48 contiguous bytes) and x (16 B) +# by 16-byte loads, then the TRT-LLM kernel's arithmetic instruction for instruction +# (`__fmaf_rn` over the taps from 0, `v * __frcp_rn(1 + __expf(-v))`, round to bf16), +# x written in place and the state shifted by the new token. +# -------------------------------------------------------------------------------------- + +#: Rows from which the decode shape runs the CuTe kernel (GB300, graph replay, 8 +#: rotating layers, us: 128 rows 7.46 -> 5.41, 121 rows 7.40 -> 5.13; at 33 rows +#: and below the TRT-LLM kernel's one wave is faster: 3.10 vs 3.39). +_CUTE_MIN_ROWS = 64 +_CH = 8 +_THREADS = 128 + + +def _iv(x, loc, ip): + return x.ir_value(loc=loc, ip=ip) + + +def _asm(ret, args, text, constraints, loc, ip, side_effects=False): + return llvm.inline_asm( + ret, + [_iv(a, loc, ip) for a in args], + text, + constraints, + has_side_effects=side_effects, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _asm4(args, text, constraints, loc, ip, side_effects=False): + ty = T.i32() + res = _asm( + llvm.StructType.get_literal([ty] * 4), args, text, constraints, loc, ip, side_effects + ) + return [cutlass.Int32(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip)) for i in range(4)] + + +@dsl_user_op +def _ldg_v4(addr, *, loc=None, ip=None): + """16 bytes; volatile so that no load moves across the pool's in-place update.""" + return _asm4( + [cutlass.Int64(addr)], + "ld.global.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _ldg_nc_v4(addr, *, loc=None, ip=None): + """16 bytes of a read-only operand (the weights), non-coherent path.""" + return _asm4( + [cutlass.Int64(addr)], + "ld.global.nc.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + loc, + ip, + ) + + +@dsl_user_op +def _stg_v4(addr, a, b, c, d, *, loc=None, ip=None): + _asm( + T.i32(), + [cutlass.Int64(addr)] + [cutlass.Int32(v) for v in (a, b, c, d)], + "{ st.global.v4.b32 [$1], {$2, $3, $4, $5}; mov.u32 $0, 0; }", + "=r,l,r,r,r,r", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _bf16_lo(w, *, loc=None, ip=None): + """The low bf16 of a word as fp32 (exact).""" + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Int32(w)], + "{ .reg .b32 t; shl.b32 t, $1, 16; mov.b32 $0, t; }", + "=f,r", + loc, + ip, + ) + ) + + +@dsl_user_op +def _bf16_hi(w, *, loc=None, ip=None): + """The high bf16 of a word as fp32 (exact).""" + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Int32(w)], + "{ .reg .b32 t; and.b32 t, $1, 0xFFFF0000; mov.b32 $0, t; }", + "=f,r", + loc, + ip, + ) + ) + + +@dsl_user_op +def _bf16x2(lo, hi, *, loc=None, ip=None): + """`cvt.rn.bf16x2.f32`: `lo` in the low half (exact for values already in bf16).""" + return cutlass.Int32( + _asm( + T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "cvt.rn.bf16x2.f32 $0, $1, $2;", + "=r,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _conv_silu(w0, w1, w2, w3, s0, s1, s2, x, *, loc=None, ip=None): + """`__fmaf_rn(w[i], s[i], v)` over the four taps from 0, then `v * __frcp_rn(1 + + __expf(-v))` (`__expf` is `ex2.approx.ftz` of the `mul.ftz` by log2e).""" + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Float32(v) for v in (w0, w1, w2, w3, s0, s1, s2, x)], + "{ .reg .f32 v, t, e; fma.rn.f32 v, $1, $5, 0f00000000; fma.rn.f32 v, $2, $6, v; " + "fma.rn.f32 v, $3, $7, v; fma.rn.f32 v, $4, $8, v; neg.f32 t, v; " + "mul.ftz.f32 t, t, 0f3FB8AA3B; ex2.approx.ftz.f32 e, t; add.f32 e, e, 0f3F800000; " + "rcp.rn.f32 e, e; mul.rn.f32 $0, v, e; }", + "=f,f,f,f,f,f,f,f,f", + loc, + ip, + ) + ) + + +class _DecodeConvKernel: + """One token a row: `x[row, c] = silu(conv)` in place and the row's slot shifted.""" + + def __init__(self, dim: int, x_stride: int): + assert dim % (_CH * _THREADS) == 0 + self.dim, self.sx = dim, x_stride + + @cute.jit + def __call__( + self, + x_t: cute.Tensor, + w_ptr: cute.Pointer, + st_ptr: cute.Pointer, + slot_ptr: cute.Pointer, + rows: cutlass.Int32, + slot_stride: cutlass.Int32, + pad_slot_id: cutlass.Int32, + stream, + ): + slots = cute.make_tensor(slot_ptr, cute.make_layout((rows,), stride=(1,))) + # `x` arrives as a tensor: the TVM-FFI entry takes its stream from a GPU tensor. + self.kernel( + x_t.iterator.toint(), w_ptr.toint(), st_ptr.toint(), slots, slot_stride, pad_slot_id + ).launch( + grid=[rows, self.dim // (_CH * _THREADS), 1], block=[_THREADS, 1, 1], stream=stream + ) + + @cute.kernel + def kernel( + self, + x_base: cutlass.Int64, + w_base: cutlass.Int64, + st_base: cutlass.Int64, + mSlot: cute.Tensor, + slot_stride: cutlass.Int32, + pad_slot_id: cutlass.Int32, + ): + tidx, _, _ = cute.arch.thread_idx() + row, cb, _ = cute.arch.block_idx() + c0 = (cb * _THREADS + tidx) * _CH + # The weights: 8 channels x 4 taps, 32 contiguous bf16. + wv = [] + for q in cutlass.range_constexpr(4): + wv += _ldg_nc_v4(w_base + cutlass.Int64(c0 * 4 + q * 8) * 2) + slot = mSlot[row] + if slot != pad_slot_id: + sb = st_base + (cutlass.Int64(slot) * slot_stride + c0 * 3) * 2 + sv = [] + for q in cutlass.range_constexpr(3): + sv += _ldg_v4(sb + q * 16) + xa = x_base + (cutlass.Int64(row) * self.sx + c0) * 2 + xv = _ldg_v4(xa) + + def elem(words, e): + return _bf16_lo(words[e // 2]) if e % 2 == 0 else _bf16_hi(words[e // 2]) + + outs, shifted = [], [] + for ch in cutlass.range_constexpr(_CH): + w = [elem(wv, ch * 4 + k) for k in range(4)] + st = [elem(sv, ch * 3 + k) for k in range(3)] + xc = elem(xv, ch) + outs.append(_conv_silu(*w, *st, xc)) + shifted += [st[1], st[2], xc] + _stg_v4(xa, *[_bf16x2(outs[2 * i], outs[2 * i + 1]) for i in range(4)]) + for q in cutlass.range_constexpr(3): + _stg_v4( + sb + q * 16, + *[ + _bf16x2(shifted[8 * q + 2 * i], shifted[8 * q + 2 * i + 1]) + for i in range(4) + ], + ) + + +_DECODE_CONV: dict = {} + + +def _decode_conv_ok(x, conv_state, weight, bias, silu_activation, cache_seqlens, indices) -> bool: + """The decode shape this kernel certifies; anything else runs the TRT-LLM op.""" + return ( + x.dim() == 2 + and x.shape[0] >= _CUTE_MIN_ROWS + and bias is None + and silu_activation + and cache_seqlens is None + and indices is not None + and indices.dtype == torch.int32 + and indices.is_contiguous() + and indices.shape == (x.shape[0],) + and x.dtype == conv_state.dtype == weight.dtype == torch.bfloat16 + and weight.shape == (x.shape[1], 4) + and weight.is_contiguous() + and conv_state.dim() == 3 + and conv_state.shape[1:] == (x.shape[1], 3) + and conv_state.stride()[1:] == (3, 1) + and x.shape[1] % (_CH * _THREADS) == 0 + and x.stride(1) == 1 + and (x.stride(0) * 2) % 16 == 0 + and (conv_state.stride(0) * 2) % 16 == 0 + and x.data_ptr() % 16 == 0 + and conv_state.data_ptr() % 16 == 0 + and weight.data_ptr() % 16 == 0 + ) + + +def _decode_conv(x, conv_state, weight, indices, pad_slot_id) -> None: + dim = x.shape[1] + key = (dim, x.stride(0)) + args = (x, weight.data_ptr(), conv_state.data_ptr(), indices.data_ptr()) + fn = _DECODE_CONV.get(key) + if fn is None: + gm = cute.AddressSpace.gmem + cargs = ( + from_dlpack(x, assumed_align=16).mark_layout_dynamic(leading_dim=1), + make_ptr(cutlass.BFloat16, args[1], gm, assumed_align=16), + make_ptr(cutlass.BFloat16, args[2], gm, assumed_align=16), + make_ptr(cutlass.Int32, args[3], gm, assumed_align=4), + ) + fn = cute.compile( + _DecodeConvKernel(dim, x.stride(0)), + *cargs, + cutlass.Int32(1), + cutlass.Int32(1), + cutlass.Int32(-1), + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--opt-level 2 --enable-tvm-ffi", + ) + _DECODE_CONV[key] = fn + fn(*args, x.shape[0], conv_state.stride(0), pad_slot_id) + + def _pinned(batch: int) -> tuple[int, ...]: """`batch` distinct slots of the engine pool, slot 0 and the last pinned in.""" if batch == 1: @@ -51,7 +340,8 @@ class CausalConv1dUpdate(OpWrapper): # 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"), + # pad_slot_id, x_form ("3d" | "2d" | "2d_row_slice" | "2d_in_proj" | + # "token_major"), # strided_state. CELLS: tuple[Cell, ...] = ( tuple( @@ -79,6 +369,23 @@ class CausalConv1dUpdate(OpWrapper): ), ), ) + + tuple( + Cell( + why=f"x the q|k|v column view of the [rows, 16480] in_proj, batch {batch}", + spec=dict( + dim=CKPT_DIM, + batch=batch, + pool_slots=pool, + slots=slots, + x_form="2d_in_proj", + ), + ) + for batch, pool, slots in ( + (1, ENGINE_SLOTS, _pinned(1)), + (32, ENGINE_SLOTS, _pinned(32)), + (128, 256, "random"), + ) + ) + tuple( Cell( why=f"decode batch {batch}, [batch, dim, 1] x, 128-slot pool", @@ -179,11 +486,19 @@ class CausalConv1dUpdate(OpWrapper): 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 + Call site: the GDN decode step -- 2-D `[num_decodes, 10240]` x, the q|k|v + column view (row stride 16480) of the decode rows of the layer's one + `[T, 16480]` input-projection output (the op writes those columns only); + the layer's whole conv pool `[slots, 10240, 3]`, W 4, no bias, SiLU on, + `cache_seqlens` None, int32 slots, `pad_slot_id` -1. + That exact shape (silu, no ring, `state_len == W - 1`) runs TRT-LLM's + dedicated seqlen-1 kernel; every other cell runs its general one. From + `_CUTE_MIN_ROWS` (64) rows, with bf16 operands, a contiguous weight, a + pool whose slot is `[dim, 3]` contiguous and 16-byte aligned rows, it runs + `_DecodeConvKernel` (CuTe DSL) instead: 8 channels a thread by 16-byte + loads, the TRT-LLM kernel's instructions, bitwise its result (x and pool, + pad rows untouched); 128 rows 7.46 -> 5.41 us a layer (GB300, graph + replay over 8 rotating layers). Certified at 128 and 461 pool slots; at 461 the result is bitwise identical to 128 for the same slot content. @@ -216,6 +531,11 @@ def __call__( conv_state_indices: Optional[torch.Tensor], pad_slot_id: int, ) -> torch.Tensor: + if _decode_conv_ok( + x, conv_state, weight, bias, silu_activation, cache_seqlens, conv_state_indices + ): + _decode_conv(x, conv_state, weight, conv_state_indices, pad_slot_id) + return x torch.ops.trtllm.causal_conv1d_update( x.unsqueeze(-1) if x.dim() == 2 else x, conv_state, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/extract_transpose_prefill_slice.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/extract_transpose_prefill_slice.py index b2f09209e898..938b668b5d7b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/extract_transpose_prefill_slice.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/extract_transpose_prefill_slice.py @@ -20,8 +20,10 @@ ) #: The GDN mixer's packed `[q|k|v]` width for this checkpoint; the target hands -#: the whole row to this op. +#: this op the `[q|k|v]` column view of its one input-projection output. GDN_MIXED_QKV_WIDTH = 10240 +#: That output's width: `[q|k|v | z | a | b]` = 10240 + 6144 + 48 + 48. +GDN_IN_PROJ_WIDTH = 16480 class ExtractTransposePrefillSlice(OpWrapper): @@ -89,6 +91,21 @@ class ExtractTransposePrefillSlice(OpWrapper): dtype=torch.bfloat16, ), ), + *( + Cell( + why=f"GDN prefill, the q|k|v column view of the in_proj, {n} tokens", + spec=dict( + num_rows=n, + num_cols=GDN_IN_PROJ_WIDTH, + col_slice=(0, GDN_MIXED_QKV_WIDTH), + n=n, + start_col=0, + width=GDN_MIXED_QKV_WIDTH, + dtype=torch.bfloat16, + ), + ) + for n in (33, 8192) + ), Cell( why="column-slice view: row stride wider than the extracted width", spec=dict( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_post_conv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_post_conv.py index b97bbdec86f4..c83653169fe6 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_post_conv.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_post_conv.py @@ -66,7 +66,9 @@ class FusedGdnPostConv(OpWrapper): # (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)); + # "prefill_view" (prefill a `.t()` of a wider buffer, stride (1, 16384)) | + # "in_proj" (the target's one [T, 16480] input projection: decode the + # q|k|v column view of its decode rows, a/b its column views); # beta_dtype, gate_dtype (A_log/dt_bias), dt_bias_fill. CELLS: tuple[Cell, ...] = ( tuple( @@ -85,6 +87,13 @@ class FusedGdnPostConv(OpWrapper): why="target mixed call: decode a row slice of the projection", spec=dict(num_prefill=200, num_decode=32, layout="target_mixed"), ), + *( + Cell( + why=f"target in_proj views, {p} prefill + {d} decode tokens", + spec=dict(num_prefill=p, num_decode=d, layout="in_proj"), + ) + for p, d in ((200, 32), (2048, 0), (0, 128)) + ), Cell( why="stock mixer views: strided decode, a/b column halves", spec=dict(num_prefill=96, num_decode=12, layout="production"), @@ -124,8 +133,9 @@ class FusedGdnPostConv(OpWrapper): `(mq_p_t, mq_d, a, b, ...)` and the MTP-verify context half `(mq_p_t, None, a_p, b_p, ...)`, all with `l2_norm_eps=1e-6`, `softplus_threshold=20.0` and the fp32 `beta_dtype` default; `prefill` is - `extract_transpose_prefill_slice`'s contiguous output, `a`/`b` the - `[T, 48]` projections. The stock mixer's verify call passes + `extract_transpose_prefill_slice`'s contiguous output; `decode` and `a`/`b` + column views (row stride 16480) of the layer's one `[T, 16480]` + input-projection output. The stock mixer's verify call passes `beta_dtype=b.dtype`. Preconditions: `prefill.shape[0]` and `decode.shape[1]` >= qkv_dim (extra diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_sigmoid_gating_delta_rule_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_sigmoid_gating_delta_rule_update.py index c822cc138903..e487455c7ebb 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_sigmoid_gating_delta_rule_update.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_sigmoid_gating_delta_rule_update.py @@ -65,8 +65,9 @@ class FusedSigmoidGatingDeltaRuleUpdate(OpWrapper): # 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": + # halves of a [T, 96] projection | "separate": two contiguous [T, 48] | + # "in_proj": q/k/v/a/b all column views of one [T, 16480] projection, as the + # target's one input GEMM makes them), layout ("varlen" | "layer_view": # pool[1] of a [3, slots, ...] cache | "batched": cu_seqlens None, # contiguous copies | "rows": cu_seqlens None, one token per row, q/k/v # [N, 1, heads, 128] views of the packed [N, 10240] block whose token @@ -81,6 +82,17 @@ class FusedSigmoidGatingDeltaRuleUpdate(OpWrapper): ) for n in (1, 4, 16, 32) ) + + tuple( + Cell( + why=f"target decode call, in_proj column views, N={n}", + spec=dict(seq_lens=(1,) * n, slots=slots, indices=indices, ab="in_proj"), + ) + for n, slots, indices in ( + (1, ENGINE_SLOTS, _scattered(1)), + (32, ENGINE_SLOTS, _scattered(32)), + (128, 128, None), + ) + ) + ( Cell( why="varlen [3, 5, 1, 7] at the 461-slot pool", @@ -215,9 +227,10 @@ class FusedSigmoidGatingDeltaRuleUpdate(OpWrapper): layer's whole pool `[slots, 48, 128, 128]` at `mamba_ssm_cache_dtype`: - fp32 pool: `(a_log, a, dt_bias, 1.0, 20.0, q, k, v, b, ssm_states, si_all, None, True, qsl_long, None)` -- q/k/v `[1, T, heads, 128]` - reshaped column views of the conv output `[T, 10240]` (read through - `stride(1)`, no repack), `a`/`b` two contiguous `[T, 48]` projections, - int64 `cu_seqlens`; + reshaped column views of the conv output (read through `stride(1)`, no + repack), and `a`/`b`, all column views of the layer's one `[T, 16480]` + input-projection output (row stride 16480; the conv wrote q|k|v in + place), int64 `cu_seqlens`; - bf16 pool: the rows form, `cu_seqlens` None -- q/k/v `[T, 1, heads, 128]` split from the whole `[T, 10240]` block reshaped to `[T, 1, 80, 128]`, so the token stride is the row stride (`is_valid` refuses other diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_gated_norm_pdl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_gated_norm_pdl.py new file mode 100644 index 000000000000..1686f08195a9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_gated_norm_pdl.py @@ -0,0 +1,504 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One Gated-DeltaNet decode step per row, and the gated RMS norm on its output, in +one PDL launch. + +Tier 5 (CuTe DSL written for this target; the recurrence is the gated delta rule +as `gated_delta_rule_decode_pretranspose` states it, reimplemented here). At one +decode row the step reads and writes each v-head's 64 KB fp32 state once and +otherwise does a few hundred flops per state row; the rest of its time is latency. +The state is not produced by anything in the step ahead of this kernel -- the +pool row is this layer's own and was last written by this kernel one step ago, and +the slot index is filled before the step's graph is launched -- so it is read +before `griddepcontrol.wait`, while the input projection ahead of this kernel +streams its weight. + +A v-head is split over a cluster of `CLUSTER` CTAs of `WARPS` warps; each warp owns +`V / (CLUSTER * WARPS)` value rows of the head's `[V, K]` state and lane `l` the +`K / 32` key columns from `l * K / 32` of each. Small CTAs fit beside the input +projection's last wave, so their state reads start before it ends. The gated RMS +norm (per v-head, 128 wide) needs the head's whole output: every warp writes the +sum of squares of its rows into each CTA of the cluster through distributed shared +memory, and every CTA then normalizes and gates its own rows. +""" + +from __future__ import annotations + +import functools +import os + +import cuda.bindings.driver as cuda_drv +import cutlass +import cutlass.cute as cute +import torch +from cutlass import const_expr +from cutlass._mlir.dialects import llvm +from cutlass.cute import experimental as cute_ext +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import dsl_user_op + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.gated_delta_rule_decode_pretranspose import ( # noqa: E501 + gated_delta_rule_decode_pretranspose as _recurrence, +) + +from ..norm.rms_norm_gated_token_major import rms_norm_gated_token_major as _gated_norm + +#: Decode rows per launch this entry is certified for. +MAX_ROWS = 4 +#: CTAs per v-head (one cluster) and warps per CTA; their product must be <= 32. +CLUSTER = 8 +WARPS = 4 + + +@functools.lru_cache(maxsize=None) +def pdl_enabled(device_index: int) -> bool: + """Programmatic dependent launch: sm_90+ unless TRTLLM_ENABLE_PDL=0.""" + props = torch.cuda.get_device_properties(device_index) + return props.major >= 9 and os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" + + +@dsl_user_op +def _st_shared_cluster_f32(ptr, val, *, loc=None, ip=None): + """`st.shared::cluster.f32` to a pointer from `cute.arch.map_dsmem_ptr`.""" + llvm.inline_asm( + None, + [ + cutlass.Int32(ptr.toint(loc=loc, ip=ip)).ir_value(loc=loc, ip=ip), + cutlass.Float32(val).ir_value(loc=loc, ip=ip), + ], + "st.shared::cluster.f32 [$0], $1;", + "r,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +class _GdnKernel: + def __init__(self, *, k_heads, v_heads, kdim, vdim, pdl): + assert vdim % (CLUSTER * WARPS) == 0 and CLUSTER * WARPS <= 32 + assert kdim % 32 == 0 and v_heads % k_heads == 0 + self.k_heads, self.v_heads, self.kdim, self.vdim, self.pdl = ( + k_heads, + v_heads, + kdim, + vdim, + pdl, + ) + self.rows_per_cta = vdim // CLUSTER + self.rows_per_warp = vdim // (CLUSTER * WARPS) + self.per_lane = kdim // 32 + self.k_off = k_heads * kdim + self.v_off = 2 * k_heads * kdim + self.z_off = self.v_off + v_heads * vdim + self.a_off = self.z_off + v_heads * vdim + self.b_off = self.a_off + v_heads + + @cute.jit + def __call__( + self, + mProj, + mPool, + mSlot, + mALog, + mDt, + mW, + mOut, + eps: cutlass.Float32, + stream: cuda_drv.CUstream, + ): + rows = cute.size(mProj, mode=[0]) + ld_state = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + cutlass.Float32, + num_bits_per_copy=32 * self.per_lane, + load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING, + ) + st_state = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=32 * self.per_lane + ) + ld_act = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + cutlass.BFloat16, + num_bits_per_copy=16 * self.per_lane, + load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS, + ) + self.kernel( + mProj, mPool, mSlot, mALog, mDt, mW, mOut, eps, ld_state, st_state, ld_act + ).launch( + grid=[self.v_heads * CLUSTER, rows, 1], + block=[WARPS * 32, 1, 1], + cluster=[CLUSTER, 1, 1], + smem=32 * 4 + 16, + stream=stream, + use_pdl=self.pdl, + # <= 64 registers a thread: the CTAs fit beside the producer's last wave. + min_blocks_per_mp=max(1, 8 * 4 // WARPS), + ) + + @cute.kernel + def kernel( + self, + mProj, + mPool, + mSlot, + mALog, + mDt, + mW, + mOut, + eps: cutlass.Float32, + ld_state: cute.CopyAtom, + st_state: cute.CopyAtom, + ld_act: cute.CopyAtom, + ): + bidx, n, _ = cute.arch.block_idx() + warp = cute.arch.warp_idx() + lane = cute.arch.lane_idx() + rpw: cutlass.Constexpr = self.rows_per_warp + pl: cutlass.Constexpr = self.per_lane + kdim: cutlass.Constexpr = self.kdim + vdim: cutlass.Constexpr = self.vdim + parts: cutlass.Constexpr = CLUSTER * WARPS + hv = bidx // CLUSTER + part = bidx % CLUSTER # the CTA's rank in its cluster: clusters tile the grid's x + kh = hv // (self.v_heads // self.k_heads) + v0 = part * self.rows_per_cta + warp * rpw + k0 = lane * pl + ss_parts = cutlass.utils.SmemAllocator().allocate_tensor( + cutlass.Float32, cute.make_layout((32,), stride=(1,)), byte_alignment=16 + ) + if const_expr(CLUSTER > 1): + # Every CTA of the cluster must have started before another writes into its + # shared memory; this arrival is waited on just before those writes. + cute.arch.cluster_arrive_relaxed() + + # 1. Static and step-invariant operands, before the wait: the gates' constants, + # the norm weight of the warp's rows, the slot and the warp's state rows. + a_log_h = mALog[hv] + dtb = mDt[hv] + wn = cute.make_rmem_tensor(cute.make_layout((rpw,), stride=(1,)), cutlass.Float32) + for r in cutlass.range_constexpr(rpw): + wn[r] = mW[v0 + r].to(cutlass.Float32) + slot = mSlot[n] + # 64-bit: the cache manager coalesces every GDN layer's state of a slot into one + # page, so a layer's slot stride is ~28M elements and slot * stride overflows int32. + slot64 = cutlass.Int64(slot) + # (slot, hv, v, (pl, lane)): lane l's 16-byte vector of a state row. + pool_vec = cute.logical_divide(mPool, (None, None, None, pl)) + h = cute.make_rmem_tensor(cute.make_layout((rpw, pl), stride=(pl, 1)), cutlass.Float32) + h.fill(0.0) + if slot >= 0: + for r in cutlass.range_constexpr(rpw): + cute.copy(ld_state, pool_vec[slot64, hv, v0 + r, (None, lane)], h[r, None]) + if const_expr(self.pdl): + cute.arch.griddepcontrol_launch_dependents() + cute.arch.griddepcontrol_wait() + + # 2. The projection's q, k, v, z, a, b for this row and head. + # (n, (pl, cols / pl)): 8-byte vectors of the row; every offset below is a multiple of pl. + proj_vec = cute.logical_divide(mProj, (None, pl)) + raw = cute.make_rmem_tensor(cute.make_layout((2, pl), stride=(pl, 1)), cutlass.BFloat16) + cute.copy(ld_act, proj_vec[n, (None, (kh * kdim + k0) // pl)], raw[0, None]) + cute.copy(ld_act, proj_vec[n, (None, (self.k_off + kh * kdim + k0) // pl)], raw[1, None]) + qn = cute.make_rmem_tensor(cute.make_layout((pl,), stride=(1,)), cutlass.Float32) + kn = cute.make_rmem_tensor(cute.make_layout((pl,), stride=(1,)), cutlass.Float32) + zv = cute.make_rmem_tensor(cute.make_layout((rpw,), stride=(1,)), cutlass.Float32) + vv = cute.make_rmem_tensor(cute.make_layout((rpw,), stride=(1,)), cutlass.Float32) + for e in cutlass.range_constexpr(pl): + qn[e] = raw[0, e].to(cutlass.Float32) + kn[e] = raw[1, e].to(cutlass.Float32) + for r in cutlass.range_constexpr(rpw): + vv[r] = mProj[n, self.v_off + hv * vdim + v0 + r].to(cutlass.Float32) + zv[r] = mProj[n, self.z_off + hv * vdim + v0 + r].to(cutlass.Float32) + a = mProj[n, self.a_off + hv].to(cutlass.Float32) + b = mProj[n, self.b_off + hv].to(cutlass.Float32) + + # q, k L2 norm (eps 1e-6 inside the root), the 1/sqrt(K) scale on q, the gates. + sq = cutlass.Float32(0.0) + sk = cutlass.Float32(0.0) + for e in cutlass.range_constexpr(pl): + sq = qn[e] * qn[e] + sq + sk = kn[e] * kn[e] + sk + sq = cute.arch.warp_reduction_sum(sq) + sk = cute.arch.warp_reduction_sum(sk) + rq = cute.math.rsqrt(sq + 1e-6) * (1.0 / cute.math.sqrt(cutlass.Float32(kdim))) + rk = cute.math.rsqrt(sk + 1e-6) + for e in cutlass.range_constexpr(pl): + qn[e] = qn[e] * rq + kn[e] = kn[e] * rk + x = a + dtb + sp = x + if x <= 20.0: + sp = cute.math.log1p(cute.math.exp(x)) + decay = cute.math.exp(-cute.math.exp(a_log_h) * sp) + beta = 1.0 / (1.0 + cute.math.exp(-b)) + + # 3. The delta rule, one value row at a time; the new rows go straight back. + ov = cute.make_rmem_tensor(cute.make_layout((rpw,), stride=(1,)), cutlass.Float32) + ssw = cutlass.Float32(0.0) + for r in cutlass.range_constexpr(rpw): + hk = cutlass.Float32(0.0) + for e in cutlass.range_constexpr(pl): + hk = h[r, e] * kn[e] + hk + hk = cute.arch.warp_reduction_sum(hk) + u = beta * (vv[r] - decay * hk) + o = cutlass.Float32(0.0) + for e in cutlass.range_constexpr(pl): + h[r, e] = decay * h[r, e] + u * kn[e] + o = h[r, e] * qn[e] + o + o = cute.arch.warp_reduction_sum(o) + if slot >= 0: + cute.copy(st_state, h[r, None], pool_vec[slot64, hv, v0 + r, (None, lane)]) + else: + o = cutlass.Float32(0.0) + # The recurrence's output is a bf16 tensor in the unfused composition. + ov[r] = o.to(cutlass.BFloat16).to(cutlass.Float32) + ssw = ov[r] * ov[r] + ssw + + # 4. The head's sum of squares: warp `w` of CTA `c` writes its rows' share into + # slot `c * WARPS + w` of every CTA of the cluster (lane `i` into CTA `i`). + if const_expr(CLUSTER > 1): + cute.arch.cluster_wait() + if lane < CLUSTER: + _st_shared_cluster_f32( + cute.arch.map_dsmem_ptr(ss_parts.iterator + (part * WARPS + warp), lane), ssw + ) + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + else: + if lane == 0: + ss_parts[warp] = ssw + cute.arch.sync_threads() + ss = cutlass.Float32(0.0) + if lane < parts: + ss = ss_parts[lane] + ss = cute.arch.warp_reduction_sum(ss) + rstd = cute.math.rsqrt(ss / vdim + eps) + + # 5. Gated RMS norm of the warp's rows: lane `r` writes row `r`. + o_r = ov[0] + z_r = zv[0] + w_r = wn[0] + for r in cutlass.range_constexpr(1, rpw): + o_r = cutlass.Float32(cutlass.select_(lane == r, ov[r], o_r)) + z_r = cutlass.Float32(cutlass.select_(lane == r, zv[r], z_r)) + w_r = cutlass.Float32(cutlass.select_(lane == r, wn[r], w_r)) + y = o_r * rstd * w_r * (1.0 / (1.0 + cute.math.exp(-z_r))) + if lane < rpw: + mOut[n, hv * vdim + v0 + lane] = y.to(cutlass.BFloat16) + + +def _dl(t: torch.Tensor, align: int = 4): + return from_dlpack(t, assumed_align=align) + + +@functools.lru_cache(maxsize=None) +def _compiled(rows, k_heads, v_heads, kdim, vdim, proj_stride, pool_shape, pool_stride, pdl): + kern = _GdnKernel(k_heads=k_heads, v_heads=v_heads, kdim=kdim, vdim=vdim, pdl=pdl) + dev = "cuda" + width = 2 * k_heads * kdim + 2 * v_heads * vdim + 2 * v_heads + proj = torch.empty(rows, proj_stride, dtype=torch.bfloat16, device=dev)[:, :width] + pool = torch.empty_strided(pool_shape, pool_stride, dtype=torch.float32, device=dev) + slot = torch.empty(rows, dtype=torch.int32, device=dev) + a_log_h = torch.empty(v_heads, dtype=torch.float32, device=dev) + dt = torch.empty(v_heads, dtype=torch.float32, device=dev) + w = torch.empty(vdim, dtype=torch.bfloat16, device=dev) + out = torch.empty(rows, v_heads * vdim, dtype=torch.bfloat16, device=dev) + stream = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream) + return cute_ext.compile( + kern, + _dl(proj, 8), + _dl(pool, 16), + _dl(slot), + _dl(a_log_h), + _dl(dt), + _dl(w, 2), + _dl(out, 2), + cutlass.Float32(1e-6), + stream, + ) + + +def launch(proj, pool, slots, A_log, dt_bias, norm_weight, eps, k_heads, kdim, pdl=True): + """The launch behind `GdnDecodeGatedNormPdl.__call__`; updates `pool` in place and + returns the gated-normed output `[rows, HV * V]` bf16.""" + rows = proj.shape[0] + v_heads, vdim = pool.shape[1], pool.shape[2] + out = torch.empty((rows, v_heads * vdim), dtype=torch.bfloat16, device=proj.device) + if rows == 0: + return out + dev = proj.device + pdl = pdl and pdl_enabled(dev.index if dev.index is not None else torch.cuda.current_device()) + fn = _compiled( + rows, + k_heads, + v_heads, + kdim, + vdim, + proj.stride(0), + tuple(pool.shape), + tuple(pool.stride()), + bool(pdl), + ) + stream = cuda_drv.CUstream(torch.cuda.current_stream(dev).cuda_stream) + fn( + _dl(proj, 8), + _dl(pool, 16), + _dl(slots), + _dl(A_log), + _dl(dt_bias), + _dl(norm_weight, 2), + _dl(out, 2), + cutlass.Float32(eps), + stream, + ) + return out + + +class GdnDecodeGatedNormPdl(OpWrapper): + """A single-token Gated-DeltaNet step per row over the fp32 slot pool, then the + sigmoid-gated RMS norm of its output per v-head. With `proj` the decode input + projection's output row `[q (H*K) | k (H*K) | v (HV*V) | z (HV*V) | a (HV) | b (HV)]` + (after the conv step: q, k, v are the conv's SiLU outputs):: + + o, pool = gated_delta_rule_decode_pretranspose(q, k, v, A_log, a, dt_bias, b, + use_qk_l2norm=True, pool, slots) + out = rms_norm_gated_token_major(o (bf16), z, norm_weight, eps, "sigmoid") + + Returns `out` `[rows, HV * V]` bf16 and updates `pool[slots[n]]` in place; a + negative slot skips the row's state (no read or write) and its output is zero. + `o` is rounded to bf16 between the two, as the composition stores it. + + Fusion boundary: both gates, the q/k L2 norm, the recurrence, the indexed state + read and write, the gated RMS norm. The projection and the conv are the caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="qwen3.8 decode, 1 row (con 1), 461-slot pool", spec=dict(rows=1)), + Cell(why="qwen3.8 decode, 4 rows (MAX_ROWS), scattered slots", spec=dict(rows=4)), + Cell(why="qwen3.8 decode, 2 rows, one pad slot (-1)", spec=dict(rows=2, pad=True)), + Cell( + why="qwen3.8 decode, 1 row, the pool a layer view of a 3-layer cache", + spec=dict(rows=1, pool="layer_view"), + ), + Cell( + why="2 rows, the cache manager's coalesced pool: slot stride 36 layers, slot 100", + spec=dict(rows=2, pool="coalesced"), + ), + ) + + note = """ + Metadata consumed: `slots` (the cache manager's state indices for the step). + + Preconditions: bf16 `proj` `[rows, 2*H*K + 2*HV*V + 2*HV]` with unit column stride; + the pool fp32 `[slots, HV, V, K]` with K innermost (any outer strides); `A_log`, + `dt_bias` fp32 `[HV]`; `norm_weight` bf16 `[V]`; `V % 32 == 0` (rows a warp: + `V / 32`), `K % 32 == 0`; int32 `slots` `[rows]`. **The pool rows and `slots` must not be + written by any kernel of the same step queued ahead of this one**: both are read + before `griddepcontrol.wait`. That holds for this target's decode steps (each + layer's pool is written only by its own recurrence, `slots` is filled before the + step's graph launches) and is why the entry is not for verify or prefill steps. + + Launch: `HV * 8` CTAs a row, clusters of 8 (one v-head), 128 threads, <= 64 + registers. Numerics: fp32 inside; `o` rounded to bf16 before the norm, the output + rounded once. Inside the recurrence entry's band (torch's default for bf16 / fp32) of the + fp64 composition. Deterministic. Graph-capturable after the first call per + geometry (compiled then). + """ + + def __call__( + self, + proj: torch.Tensor, + pool: torch.Tensor, + slots: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + k_heads: int, + kdim: int, + pdl: bool = True, + ) -> torch.Tensor: + return launch(proj, pool, slots, A_log, dt_bias, norm_weight, eps, k_heads, kdim, pdl) + + def reference( + self, + proj: torch.Tensor, + pool: torch.Tensor, + slots: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + k_heads: int, + kdim: int, + pdl: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + """`(out, pool_after)`: the recurrence entry's fp64 reference (o rounded to + bf16, the pool to fp32), then the gated norm entry's fp32 reference.""" + rows = proj.shape[0] + v_heads, vdim = pool.shape[1], pool.shape[2] + q, k, v, z, a, b = torch.split( + proj, + [k_heads * kdim, k_heads * kdim, v_heads * vdim, v_heads * vdim, v_heads, v_heads], + 1, + ) + o, pool_after = _recurrence.reference( + q=q.reshape(rows, 1, k_heads, kdim), + k=k.reshape(rows, 1, k_heads, kdim), + v=v.reshape(rows, 1, v_heads, vdim), + state=None, + A_log=A_log, + a=a.reshape(rows, 1, v_heads), + dt_bias=dt_bias, + b=b.reshape(rows, 1, v_heads), + scale=None, + output=torch.empty(rows, 1, v_heads, vdim, dtype=torch.bfloat16, device=proj.device), + use_qk_l2norm=True, + initial_state=pool, + initial_state_indices=slots, + output_state_indices=None, + ) + out = _gated_norm.reference( + o.reshape(-1, vdim), z.reshape(rows, v_heads, vdim), norm_weight, eps, None, "sigmoid" + ) + return out.reshape(rows, v_heads * vdim), pool_after + + def is_valid( + self, + proj: torch.Tensor, + pool: torch.Tensor, + slots: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + k_heads: int, + kdim: int, + pdl: bool = True, + ) -> None: + rows = proj.shape[0] + v_heads, vdim, kd = pool.shape[1], pool.shape[2], pool.shape[3] + assert rows <= MAX_ROWS, f"at most {MAX_ROWS} rows" + assert proj.dtype == torch.bfloat16 and proj.stride(1) == 1, "bf16 proj, unit column stride" + assert proj.shape[1] == 2 * k_heads * kdim + 2 * v_heads * vdim + 2 * v_heads, "proj width" + assert pool.dtype == torch.float32 and pool.stride(3) == 1 and kd == kdim, ( + "fp32 pool [S, HV, V, K]" + ) + assert vdim % (CLUSTER * WARPS) == 0 and kdim % 32 == 0, "V % (CLUSTER * WARPS), K % 32" + assert v_heads % k_heads == 0, "HV a multiple of H" + assert slots.dtype == torch.int32 and slots.shape == (rows,), "int32 slots [rows]" + assert A_log.dtype == torch.float32 and dt_bias.dtype == torch.float32, ( + "fp32 gate constants" + ) + assert norm_weight.dtype == torch.bfloat16 and norm_weight.shape == (vdim,), ( + "bf16 norm weight [V]" + ) + + +gdn_decode_gated_norm_pdl = GdnDecodeGatedNormPdl() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_rows_gated_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_rows_gated_norm.py new file mode 100644 index 000000000000..33c82416a042 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_rows_gated_norm.py @@ -0,0 +1,715 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One Gated-DeltaNet decode step per row and the gated RMS norm on its output, streamed +over the fp32 state pool: the decode step at batch sizes past the cluster kernel's. + +Tier 6 (CuTe DSL written for this target; the recurrence is the gated delta rule as +`gated_delta_rule_decode_pretranspose` states it, reimplemented here). At a decode batch +the step is the pool's traffic -- each (row, v-head) state, 64 KB fp32, read and written +once -- so a CTA owns one whole state and streams it: warp w the value rows +[w * RPW, (w + 1) * RPW), lane l the key columns [4l, 4l + 4) of each, so a warp's +access is one contiguous 512-byte state row. The state is loaded before +`griddepcontrol.wait` (nothing in the step writes it ahead of this kernel). The two dot +products a value row takes over K are reduced across the warp by a transposed +reduction (RPW rows in RPW - 1 + log2(32 / RPW) shuffles, not 5 RPW), which leaves row +j's sum in lanes [j * 32 / RPW, (j + 1) * 32 / RPW); `u` goes back to every lane by +index shuffles. The CTA holds the head's whole output, so the gated RMS norm runs in +place. +""" + +from __future__ import annotations + +import math + +import cutlass +import cutlass.cute as cute +import torch +from cutlass._mlir.dialects import llvm +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.gated_delta_rule_decode_pretranspose import ( # noqa: E501 + gated_delta_rule_decode_pretranspose as _recurrence, +) +from tensorrt_llm._torch.cute_dsl_kernels.blackwell.utils import make_ptr + +from ..norm.rms_norm_gated_token_major import rms_norm_gated_token_major as _gated_norm +from .gdn_decode_gated_norm_pdl import pdl_enabled + +#: Key columns a lane (a 16-byte vector of a state row). +_PL = 4 +#: Warps a CTA; a warp owns `V // _WARPS` value rows of the head's state. +_WARPS = 8 +#: Largest recurrent-pool slot count the cache manager was seen to allocate for this +#: checkpoint; the kernel addresses a slot by its index alone, any pool size alike. +ENGINE_SLOTS = 461 + + +def _iv(x, loc, ip): + return x.ir_value(loc=loc, ip=ip) + + +def _asm(ret, args, text, constraints, loc, ip, side_effects=False): + return llvm.inline_asm( + ret, + [_iv(a, loc, ip) for a in args], + text, + constraints, + has_side_effects=side_effects, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _asm_n(n, ty, args, text, constraints, loc, ip, side_effects=False): + res = _asm( + llvm.StructType.get_literal([ty] * n), args, text, constraints, loc, ip, side_effects + ) + return [llvm.extractvalue(ty, res, [i], loc=loc, ip=ip) for i in range(n)] + + +@dsl_user_op +def _ld_state(addr, pred, *, loc=None, ip=None): + """Four fp32 of a state row (streaming), or zeros where `pred` is 0.""" + r = _asm_n( + 4, + T.f32(), + [cutlass.Int64(addr), cutlass.Int32(pred)], + "{ .reg .pred p; setp.ne.b32 p, $5, 0; mov.b32 $0, 0; mov.b32 $1, 0; mov.b32 $2, 0; " + "mov.b32 $3, 0; @p ld.global.cs.v4.f32 {$0, $1, $2, $3}, [$4]; }", + "=f,=f,=f,=f,l,r", + loc, + ip, + side_effects=True, + ) + return [cutlass.Float32(v) for v in r] + + +@dsl_user_op +def _st_state(addr, pred, a, b, c, d, *, loc=None, ip=None): + """Four fp32 of a state row (streaming) where `pred` is not 0.""" + _asm( + T.i32(), + [cutlass.Int64(addr), cutlass.Int32(pred)] + [cutlass.Float32(v) for v in (a, b, c, d)], + "{ .reg .pred p; setp.ne.b32 p, $2, 0; @p st.global.cs.v4.f32 [$1], {$3, $4, $5, $6}; " + "mov.u32 $0, 0; }", + "=r,l,r,f,f,f,f", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _ld_bf16x4(addr, *, loc=None, ip=None): + """Four bf16 (8 bytes) as fp32.""" + r = _asm_n( + 4, + T.f32(), + [cutlass.Int64(addr)], + "{ .reg .b32 w0, w1; ld.global.v2.b32 {w0, w1}, [$4]; shl.b32 $0, w0, 16; " + "and.b32 $1, w0, 0xFFFF0000; shl.b32 $2, w1, 16; and.b32 $3, w1, 0xFFFF0000; }", + "=f,=f,=f,=f,l", + loc, + ip, + side_effects=True, + ) + return [cutlass.Float32(v) for v in r] + + +@dsl_user_op +def _ld_bf16(addr, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Int64(addr)], + "{ .reg .b16 h; ld.global.b16 h, [$1]; cvt.f32.bf16 $0, h; }", + "=f,l", + loc, + ip, + side_effects=True, + ) + ) + + +@dsl_user_op +def _st_bf16(addr, pred, x, *, loc=None, ip=None): + """`x` rounded to bf16 where `pred` is not 0.""" + _asm( + T.i32(), + [cutlass.Int64(addr), cutlass.Int32(pred), cutlass.Float32(x)], + "{ .reg .pred p; .reg .b16 h; setp.ne.b32 p, $2, 0; cvt.rn.bf16.f32 h, $3; " + "@p st.global.b16 [$1], h; mov.u32 $0, 0; }", + "=r,l,r,f", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _round_bf16(x, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Float32(x)], + "{ .reg .b16 h; cvt.rn.bf16.f32 h, $1; cvt.f32.bf16 $0, h; }", + "=f,f", + loc, + ip, + ) + ) + + +def _rows_reduce(vals, lane): + """Sum each of `vals` (a lane's partial of one row each) over the warp. Each halving + level (offset 16 first) trades half the rows with the partner lane, so row j's sum + lands in lanes [j * 32 / n, (j + 1) * 32 / n); full butterflies finish those lanes.""" + vals = list(vals) + off = 16 + while len(vals) > 1: + half = len(vals) // 2 + upper = (lane & off) != 0 + kept = [] + for i in range(half): + lo, hi = vals[i], vals[i + half] + recv = cute.arch.shuffle_sync_bfly(cutlass.select_(upper, lo, hi), offset=off) + kept.append(cutlass.select_(upper, hi, lo) + recv) + vals = kept + off //= 2 + v = vals[0] + while off >= 1: + v = v + cute.arch.shuffle_sync_bfly(v, offset=off) + off //= 2 + return v + + +class _GdnRowsKernel: + """CTA (v-head, row): the row's whole state of that head. `norm`: the gated RMS norm on + the head's output; else the recurrence's output, rounded to bf16, is stored as is.""" + + def __init__(self, k_heads, v_heads, kdim, vdim, pdl, norm=True): + assert kdim == 32 * _PL and vdim % _WARPS == 0 + self.h, self.hv, self.k, self.v, self.pdl = k_heads, v_heads, kdim, vdim, pdl + self.norm = norm + self.rpw = vdim // _WARPS + assert self.rpw in (4, 8, 16, 32) + # Lanes that hold one row's reduced sums. + self.lpr = 32 // self.rpw + self.k_off = k_heads * kdim + self.v_off = 2 * k_heads * kdim + self.z_off = self.v_off + v_heads * vdim + self.a_off = self.z_off + v_heads * vdim + self.b_off = self.a_off + v_heads + + @cute.jit + def __call__( + self, + proj_t: cute.Tensor, + pool_ptr: cute.Pointer, + slot_ptr: cute.Pointer, + a_log_ptr: cute.Pointer, + dt_ptr: cute.Pointer, + w_ptr: cute.Pointer, + out_ptr: cute.Pointer, + rows: cutlass.Int32, + s_proj: cutlass.Int32, + s_slot: cutlass.Int64, + s_head: cutlass.Int32, + s_row: cutlass.Int32, + s_out: cutlass.Int32, + eps: cutlass.Float32, + stream, + ): + slots = cute.make_tensor(slot_ptr, cute.make_layout((rows,), stride=(1,))) + a_log_t = cute.make_tensor(a_log_ptr, cute.make_layout((self.hv,), stride=(1,))) + dt = cute.make_tensor(dt_ptr, cute.make_layout((self.hv,), stride=(1,))) + # `proj` arrives as a tensor: the TVM-FFI entry takes its stream from a GPU tensor. + self.kernel( + proj_t.iterator.toint(), + pool_ptr.toint(), + slots, + a_log_t, + dt, + w_ptr.toint(), + out_ptr.toint(), + s_proj, + s_slot, + s_head, + s_row, + s_out, + eps, + ).launch( + grid=[self.hv, rows, 1], + block=[32 * _WARPS, 1, 1], + smem=_WARPS * 4 + 64, + stream=stream, + use_pdl=self.pdl, + ) + + @cute.kernel + def kernel( + self, + proj_base: cutlass.Int64, + pool_base: cutlass.Int64, + mSlot: cute.Tensor, + mALog: cute.Tensor, + mDt: cute.Tensor, + w_base: cutlass.Int64, + out_base: cutlass.Int64, + s_proj: cutlass.Int32, + s_slot: cutlass.Int64, + s_head: cutlass.Int32, + s_row: cutlass.Int32, + s_out: cutlass.Int32, + eps: cutlass.Float32, + ): + tidx, _, _ = cute.arch.thread_idx() + hv, row, _ = cute.arch.block_idx() + warp = tidx // 32 + lane = tidx % 32 + kh = hv // (self.hv // self.h) + v0 = warp * self.rpw + jl = lane // self.lpr # the value row whose sums this lane ends up holding + red = cutlass.utils.SmemAllocator().allocate_tensor( + cutlass.Float32, cute.make_layout((_WARPS,), stride=(1,)), byte_alignment=16 + ) + + # 1. Before the wait: the warp's state rows (a negative slot reads none: zero state), + # the gate constants and the norm weight of this lane's row. + slot = mSlot[row] + live = cutlass.Int32(0) + if slot >= 0: + live = cutlass.Int32(1) + # 64-bit: a coalesced cache page puts every GDN layer's state of a slot in one + # page, so `slot * s_slot` passes int32. + base = ( + pool_base + + ( + cutlass.Int64(cutlass.max(slot, cutlass.Int32(0))) * s_slot + + cutlass.Int64(hv * s_head + v0 * s_row + lane * _PL) + ) + * 4 + ) + h = [_ld_state(base + cutlass.Int64(j * s_row * 4), live) for j in range(self.rpw)] + a_log = mALog[hv] + dtb = mDt[hv] + w_l = cutlass.Float32(0.0) + if cutlass.const_expr(self.norm): + w_l = _ld_bf16(w_base + cutlass.Int64(v0 + jl) * 2) + if cutlass.const_expr(self.pdl): + cute.arch.griddepcontrol_launch_dependents() + cute.arch.griddepcontrol_wait() + + # 2. The row's q and k (this lane's four columns), v and z (this lane's value + # row), a and b; the q/k L2 norms (eps 1e-6 inside the root, 1/sqrt(K) on q) and + # the gates (softplus beta 1, threshold 20). + prow = proj_base + cutlass.Int64(row) * s_proj * 2 + q4 = _ld_bf16x4(prow + cutlass.Int64(kh * self.k + lane * _PL) * 2) + k4 = _ld_bf16x4(prow + cutlass.Int64(self.k_off + kh * self.k + lane * _PL) * 2) + v_l = _ld_bf16(prow + cutlass.Int64(self.v_off + hv * self.v + v0 + jl) * 2) + z_l = cutlass.Float32(0.0) + if cutlass.const_expr(self.norm): + z_l = _ld_bf16(prow + cutlass.Int64(self.z_off + hv * self.v + v0 + jl) * 2) + a = _ld_bf16(prow + cutlass.Int64(self.a_off + hv) * 2) + b = _ld_bf16(prow + cutlass.Int64(self.b_off + hv) * 2) + sq = cutlass.Float32(0.0) + sk = cutlass.Float32(0.0) + for e in cutlass.range_constexpr(_PL): + sq = q4[e] * q4[e] + sq + sk = k4[e] * k4[e] + sk + sq = cute.arch.warp_reduction_sum(sq) + sk = cute.arch.warp_reduction_sum(sk) + rq = cute.math.rsqrt(sq + 1e-6) * (1.0 / math.sqrt(self.k)) + rk = cute.math.rsqrt(sk + 1e-6) + qn = [q4[e] * rq for e in range(_PL)] + kn = [k4[e] * rk for e in range(_PL)] + x = a + dtb + sp = x + if x <= 20.0: + sp = cute.math.log1p(cute.math.exp(x)) + decay = cute.math.exp(-cute.math.exp(a_log) * sp) + beta = 1.0 / (1.0 + cute.math.exp(-b)) + + # 3. Per value row: hk = h.k, u = beta (v - decay hk); h = decay h + u k; o = h.q. + part = [] + for j in cutlass.range_constexpr(self.rpw): + acc = cutlass.Float32(0.0) + for e in cutlass.range_constexpr(_PL): + acc = h[j][e] * kn[e] + acc + part.append(acc) + hk = _rows_reduce(part, lane) + u_l = beta * (v_l - decay * hk) + part = [] + for j in cutlass.range_constexpr(self.rpw): + u = cute.arch.shuffle_sync(u_l, j * self.lpr) + acc = cutlass.Float32(0.0) + for e in cutlass.range_constexpr(_PL): + h[j][e] = u * kn[e] + decay * h[j][e] + acc = h[j][e] * qn[e] + acc + part.append(acc) + _st_state(base + cutlass.Int64(j * s_row * 4), live, *h[j]) + o = _rows_reduce(part, lane) + + # 4. The gated RMS norm over the head's V outputs, each rounded to bf16 first (the + # recurrence's output dtype in the composition): y = o rstd w sigmoid(z). + o_b = _round_bf16(o) * cutlass.Float32(live) + own = cutlass.Int32(0) + if lane % self.lpr == 0: + own = cutlass.Int32(1) + y = o_b + if cutlass.const_expr(self.norm): + ss = cute.arch.warp_reduction_sum(o_b * o_b * cutlass.Float32(own)) + if lane == 0: + red[warp] = ss + cute.arch.sync_threads() + tot = red[0] + for i in cutlass.range_constexpr(1, _WARPS): + tot = tot + red[i] + rstd = cute.math.rsqrt(tot / self.v + eps) + y = o_b * rstd * w_l * (1.0 / (1.0 + cute.math.exp(-z_l))) + _st_bf16(out_base + (cutlass.Int64(row) * s_out + hv * self.v + v0 + jl) * 2, own, y) + + +_COMPILED: dict = {} + + +def launch(proj, pool, slots, A_log, dt_bias, norm_weight, eps, k_heads, kdim, pdl=True, out=None): + """The launch behind `GdnDecodeRowsGatedNorm.__call__` and, with `norm_weight=None`, + `GdnDecodeRows.__call__`; updates `pool` in place and returns the gated-normed output + (or the recurrence's) `[rows, HV * V]` bf16, written into `out` when given.""" + rows = proj.shape[0] + v_heads, vdim = pool.shape[1], pool.shape[2] + norm = norm_weight is not None + if out is None: + out = torch.empty((rows, v_heads * vdim), dtype=torch.bfloat16, device=proj.device) + if rows == 0: + return out + dev = proj.device + pdl = bool( + pdl and pdl_enabled(dev.index if dev.index is not None else torch.cuda.current_device()) + ) + key = (k_heads, v_heads, kdim, vdim, pdl, norm) + ptrs = ( + proj, + pool.data_ptr(), + slots.data_ptr(), + A_log.data_ptr(), + dt_bias.data_ptr(), + (norm_weight if norm else A_log).data_ptr(), + out.data_ptr(), + ) + fn = _COMPILED.get(key) + if fn is None: + gm = cute.AddressSpace.gmem + cargs = ( + from_dlpack(proj, assumed_align=8).mark_layout_dynamic(leading_dim=1), + make_ptr(cutlass.Float32, ptrs[1], gm, assumed_align=16), + make_ptr(cutlass.Int32, ptrs[2], gm, assumed_align=4), + make_ptr(cutlass.Float32, ptrs[3], gm, assumed_align=4), + make_ptr(cutlass.Float32, ptrs[4], gm, assumed_align=4), + make_ptr(cutlass.BFloat16, ptrs[5], gm, assumed_align=2), + make_ptr(cutlass.BFloat16, ptrs[6], gm, assumed_align=2), + ) + fn = cute.compile( + _GdnRowsKernel(k_heads, v_heads, kdim, vdim, pdl, norm), + *cargs, + cutlass.Int32(1), + cutlass.Int32(1), + cutlass.Int64(1), + cutlass.Int32(1), + cutlass.Int32(1), + cutlass.Int32(1), + cutlass.Float32(1e-6), + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--opt-level 2 --enable-tvm-ffi", + ) + _COMPILED[key] = fn + fn( + *ptrs, + rows, + proj.stride(0), + pool.stride(0), + pool.stride(1), + pool.stride(2), + out.stride(0), + float(eps) if norm else 0.0, + ) + return out + + +class GdnDecodeRowsGatedNorm(OpWrapper): + """A single-token Gated-DeltaNet step per row over the fp32 slot pool, then the + sigmoid-gated RMS norm of its output per v-head. With `proj` the decode input + projection's output row `[q (H*K) | k (H*K) | v (HV*V) | z (HV*V) | a (HV) | b (HV)]` + (after the conv step: q, k, v are the conv's SiLU outputs):: + + o, pool = gated_delta_rule_decode_pretranspose(q, k, v, A_log, a, dt_bias, b, + use_qk_l2norm=True, pool, slots) + out = rms_norm_gated_token_major(o (bf16), z, norm_weight, eps, "sigmoid") + + Returns `out` `[rows, HV * V]` bf16 and updates `pool[slots[n]]` in place; a + negative slot skips the row's state (no read or write) and its output is zero. + `o` is rounded to bf16 between the two, as the composition stores it. + + Fusion boundary: both gates, the q/k L2 norm, the recurrence, the indexed state + read and write, the gated RMS norm. The projection and the conv are the caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="qwen3.8 decode, 128 rows (con 128), 461-slot pool", spec=dict(rows=128)), + Cell(why="qwen3.8 decode, 3 rows (past the conv-fused path)", spec=dict(rows=3)), + Cell(why="qwen3.8 decode, 33 rows, one pad slot (-1)", spec=dict(rows=33, pad=True)), + Cell( + why="qwen3.8 decode, 16 rows, the pool a layer view of a 3-layer cache", + spec=dict(rows=16, pool="layer_view"), + ), + Cell( + why="2 rows, the cache manager's coalesced pool: slot stride 36 layers, slot 100", + spec=dict(rows=2, pool="coalesced"), + ), + ) + + note = """ + Metadata consumed: `slots` (the cache manager's state indices for the step). + + Preconditions: bf16 `proj` `[rows, 2*H*K + 2*HV*V + 2*HV]` with unit column stride + and 8-byte aligned rows; the pool fp32 `[slots, HV, V, K]` with K innermost (any + outer strides, multiples of 4) and 16-byte aligned; `K == 128`, `V % 8 == 0` with + `V / 8` in {4, 8, 16, 32}; `A_log`, `dt_bias` fp32 `[HV]`; `norm_weight` bf16 + `[V]`; int32 `slots` `[rows]`. **The pool rows and `slots` must not be written by + any kernel of the same step queued ahead of this one**: both are read before + `griddepcontrol.wait` -- as `gdn_decode_gated_norm_pdl` states, this target's decode + steps keep that (a layer's pool is written only by its own recurrence, `slots` is + filled before the step's graph launches), and verify and prefill steps do not use + the entry. + + Launch: one CTA of 8 warps a (v-head, row), `HV * rows` CTAs, PDL (off with + TRTLLM_ENABLE_PDL=0). GB300, graph replay, 461-slot pool, us (FlashInfer's + recurrence + the gated norm entry -> this): 128 rows 132.2 -> 118.8 (6.78 TB/s + against an in-place `mul_` over the same bytes at 6.83), 64 rows 70.7 -> 62.7, 16 + rows 19.1 -> 14.9, 4 rows 8.5 -> 8.4. + + Numerics: fp32 inside; each dot product over K is a lane's 4-term FMA chain and then + the transposed reduction's tree; `o` rounded to bf16 before the norm, the output + rounded once. Inside the recurrence entry's band (torch's default for bf16 / fp32) + of the composition's float64 reference, not bitwise FlashInfer's kernel. + """ + + def __call__( + self, + proj: torch.Tensor, + pool: torch.Tensor, + slots: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + k_heads: int, + kdim: int, + pdl: bool = True, + ) -> torch.Tensor: + return launch(proj, pool, slots, A_log, dt_bias, norm_weight, eps, k_heads, kdim, pdl) + + def reference( + self, + proj: torch.Tensor, + pool: torch.Tensor, + slots: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + k_heads: int, + kdim: int, + pdl: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + """`(out, pool_after)`: the recurrence entry's fp64 reference (o rounded to + bf16, the pool to fp32), then the gated norm entry's fp32 reference.""" + rows = proj.shape[0] + v_heads, vdim = pool.shape[1], pool.shape[2] + q, k, v, z, a, b = torch.split( + proj, + [k_heads * kdim, k_heads * kdim, v_heads * vdim, v_heads * vdim, v_heads, v_heads], + 1, + ) + o, pool_after = _recurrence.reference( + q=q.reshape(rows, 1, k_heads, kdim), + k=k.reshape(rows, 1, k_heads, kdim), + v=v.reshape(rows, 1, v_heads, vdim), + state=None, + A_log=A_log, + a=a.reshape(rows, 1, v_heads), + dt_bias=dt_bias, + b=b.reshape(rows, 1, v_heads), + scale=None, + output=torch.empty(rows, 1, v_heads, vdim, dtype=torch.bfloat16, device=proj.device), + use_qk_l2norm=True, + initial_state=pool, + initial_state_indices=slots, + output_state_indices=None, + ) + out = _gated_norm.reference( + o.reshape(-1, vdim), z.reshape(rows, v_heads, vdim), norm_weight, eps, None, "sigmoid" + ) + return out.reshape(rows, v_heads * vdim), pool_after + + def is_valid( + self, + proj: torch.Tensor, + pool: torch.Tensor, + slots: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + k_heads: int, + kdim: int, + pdl: bool = True, + ) -> None: + rows = proj.shape[0] + v_heads, vdim, kd = pool.shape[1], pool.shape[2], pool.shape[3] + assert proj.dtype == torch.bfloat16 and proj.stride(1) == 1, "bf16 proj, unit column stride" + assert proj.shape[1] == 2 * k_heads * kdim + 2 * v_heads * vdim + 2 * v_heads, "proj width" + assert proj.data_ptr() % 8 == 0 and proj.stride(0) % 4 == 0, "8-byte aligned proj rows" + assert pool.dtype == torch.float32 and pool.stride(3) == 1 and kd == kdim, ( + "fp32 pool [S, HV, V, K]" + ) + assert pool.data_ptr() % 16 == 0 and all(s % 4 == 0 for s in pool.stride()[:3]), ( + "16-byte aligned pool rows" + ) + assert kdim == 32 * _PL, f"K == {32 * _PL}" + assert vdim % _WARPS == 0 and vdim // _WARPS in (4, 8, 16, 32), "V / 8 in {4, 8, 16, 32}" + assert v_heads % k_heads == 0, "HV a multiple of H" + assert slots.dtype == torch.int32 and slots.shape == (rows,), "int32 slots [rows]" + assert slots.is_contiguous(), "contiguous slots" + assert A_log.dtype == torch.float32 and dt_bias.dtype == torch.float32, ( + "fp32 gate constants" + ) + assert A_log.is_contiguous() and dt_bias.is_contiguous(), "contiguous gate constants" + assert norm_weight.dtype == torch.bfloat16 and norm_weight.shape == (vdim,), ( + "bf16 norm weight [V]" + ) + + +class GdnDecodeRows(OpWrapper): + """`GdnDecodeRowsGatedNorm` without the gated norm: the recurrence's output `o` (rounded + to bf16) written into `out` `[rows, HV * V]` (unit column stride, any row stride) -- + a mixed step's generation rows, whose gated norm runs with the context rows'. The same + kernel, operands and preconditions otherwise (`proj`'s z columns are not read). + + Fusion boundary: both gates, the q/k L2 norm, the recurrence, the indexed state read and + write. Returns `out`. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="mixed step: 121 generation rows into a row range of the step output", + spec=dict(rows=121, out="row_range"), + ), + Cell(why="3 rows, one pad slot (-1)", spec=dict(rows=3, pad=True, out="row_range")), + Cell( + why="2 rows, the cache manager's coalesced pool: slot stride 36 layers", + spec=dict(rows=2, pool="coalesced", out="row_range"), + ), + ) + + note = """ + Metadata consumed: `slots`. Preconditions and launch: `GdnDecodeRowsGatedNorm`'s, plus + a bf16 `out` `[rows, HV * V]` with unit column stride. The pool rows and `slots` are + read before `griddepcontrol.wait`: in a mixed step nothing queued ahead of the call + writes the generation rows' slots (the context rows' reset and chunked recurrence touch + only theirs). Numerics: the recurrence of `GdnDecodeRowsGatedNorm`, bitwise its `o`. + """ + + def __call__( + self, + proj: torch.Tensor, + pool: torch.Tensor, + slots: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + k_heads: int, + kdim: int, + out: torch.Tensor, + pdl: bool = True, + ) -> torch.Tensor: + return launch(proj, pool, slots, A_log, dt_bias, None, 0.0, k_heads, kdim, pdl, out=out) + + def reference( + self, + proj: torch.Tensor, + pool: torch.Tensor, + slots: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + k_heads: int, + kdim: int, + out: torch.Tensor, + pdl: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + """`(o, pool_after)`: the recurrence entry's fp64 reference, o rounded to bf16.""" + rows = proj.shape[0] + v_heads, vdim = pool.shape[1], pool.shape[2] + q, k, v, _, a, b = torch.split( + proj, + [k_heads * kdim, k_heads * kdim, v_heads * vdim, v_heads * vdim, v_heads, v_heads], + 1, + ) + o, pool_after = _recurrence.reference( + q=q.reshape(rows, 1, k_heads, kdim), + k=k.reshape(rows, 1, k_heads, kdim), + v=v.reshape(rows, 1, v_heads, vdim), + state=None, + A_log=A_log, + a=a.reshape(rows, 1, v_heads), + dt_bias=dt_bias, + b=b.reshape(rows, 1, v_heads), + scale=None, + output=torch.empty(rows, 1, v_heads, vdim, dtype=torch.bfloat16, device=proj.device), + use_qk_l2norm=True, + initial_state=pool, + initial_state_indices=slots, + output_state_indices=None, + ) + return o.reshape(rows, v_heads * vdim), pool_after + + def is_valid( + self, + proj: torch.Tensor, + pool: torch.Tensor, + slots: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + k_heads: int, + kdim: int, + out: torch.Tensor, + pdl: bool = True, + ) -> None: + v_heads, vdim = pool.shape[1], pool.shape[2] + gdn_decode_rows_gated_norm.is_valid( + proj, + pool, + slots, + A_log, + dt_bias, + torch.empty(vdim, dtype=torch.bfloat16), + 0.0, + k_heads, + kdim, + ) + assert out.dtype == torch.bfloat16 and out.dim() == 2, "bf16 out [rows, HV * V]" + assert out.shape == (proj.shape[0], v_heads * vdim) and out.stride(1) == 1, ( + "out [rows, HV * V], unit column stride" + ) + + +gdn_decode_rows_gated_norm = GdnDecodeRowsGatedNorm() +gdn_decode_rows = GdnDecodeRows() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_in_proj_conv1d_pdl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_in_proj_conv1d_pdl.py new file mode 100644 index 000000000000..de6fd934aaa1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_in_proj_conv1d_pdl.py @@ -0,0 +1,443 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Gated-DeltaNet input projection at decode rows with the causal conv1d step +as its epilogue, one PDL launch. + +Tier 5 (CuTe DSL written for this target). The decode GDN block projects into one +`[T, q|k|v | z | a | b]` buffer and then runs a separate depthwise conv kernel +over the q|k|v columns. That conv touches 3 cached values and 4 taps per column: +far too little work for a launch of its own. Here every CTA computes `OUTS` +output columns over the whole K, and for a column below `conv_dim` it finishes +the conv step on the projected value before storing it. + +Order of work inside a CTA: + +1. its weight rows into registers, the conv taps of its columns, the row's slot + and the cached window of each of its conv columns -- static, host-filled, or + last written by this kernel one step ago -- so before `griddepcontrol.wait` + (PDL), while the kernel ahead of this one is still running; +2. `griddepcontrol.wait`; +3. the activation, the products, a warp then a cross-warp reduction; +4. per column: the projection rounded to bf16 (the value the unfused path + stores), and for a conv column of a non-pad row the conv in fp32 over + `[window, projected]`, SiLU, rounded once; the window advances to + `[window[1:], projected]`. + +Element `k` of a row is held by thread `(k // VEC) % THREADS` as its vector +`(k // VEC) // THREADS`, so each k-tile of loads is contiguous across a warp. +""" + +from __future__ import annotations + +import functools +import os + +import cuda.bindings.driver as cuda_drv +import cutlass +import cutlass.cute as cute +import torch +from cutlass import const_expr +from cutlass._mlir.dialects import llvm +from cutlass.cute import experimental as cute_ext +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import dsl_user_op + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +VEC = 8 +#: Before its PDL wait a CTA also has the weight rows of CTAs `bidx + j * L2_DIST` +#: (j = 1 .. L2_AHEAD) brought into L2: the kernel ahead of this one is latency-bound and +#: only a few of these CTAs fit beside it, so the later waves' weights would otherwise +#: start streaming only after it. +L2_AHEAD = 2 +L2_DIST = 1216 + + +@dsl_user_op +def _bulk_prefetch_l2(src, nbytes, *, loc=None, ip=None): + llvm.inline_asm( + None, + [ + cutlass.Int64(src.toint(loc=loc, ip=ip)).ir_value(loc=loc, ip=ip), + cutlass.Int32(nbytes).ir_value(loc=loc, ip=ip), + ], + "cp.async.bulk.prefetch.L2.global [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +#: Largest decode row count this entry serves (every row re-reads the weight from registers). +MAX_ROWS = 2 +#: (threads, outputs per CTA) per K; swept on GB300 at the checkpoint's 16480 x 2560. +TILING = {2560: (64, 2)} +WIDTH = 4 # conv taps +_COMPILE_OPTIONS = "--ptxas-options -maxrregcount=64" + + +@functools.lru_cache(maxsize=None) +def pdl_enabled(device_index: int) -> bool: + """Programmatic dependent launch: sm_90+ unless TRTLLM_ENABLE_PDL=0.""" + props = torch.cuda.get_device_properties(device_index) + return props.major >= 9 and os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" + + +class _InProjConvKernel: + def __init__(self, *, rows, k, threads, outs, conv_dim, pad_slot, pdl): + assert k % (threads * VEC) == 0 and threads % 32 == 0 + self.rows, self.threads, self.outs = rows, threads, outs + self.k_tiles = k // (threads * VEC) + self.warps = threads // 32 + self.conv_dim, self.pad_slot, self.pdl = conv_dim, pad_slot, pdl + + @cute.jit + def __call__(self, mA, mW, mOut, mState, mConvW, mSlot, stream: cuda_drv.CUstream): + n = cute.size(mW, mode=[0]) + load_a = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + cutlass.BFloat16, + num_bits_per_copy=16 * VEC, + load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS, + ) + load_w = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + cutlass.BFloat16, + num_bits_per_copy=16 * VEC, + load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING, + ) + self.kernel(mA, mW, mOut, mState, mConvW, mSlot, load_a, load_w).launch( + grid=[cute.ceil_div(n, self.outs), 1, 1], + block=[self.threads, 1, 1], + smem=self.rows * self.outs * self.warps * 4, + stream=stream, + use_pdl=self.pdl, + min_blocks_per_mp=1, + ) + + @cute.kernel + def kernel( + self, + mA, + mW, + mOut, + mState, + mConvW, + mSlot, + load_a: cute.CopyAtom, + load_w: cute.CopyAtom, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + warp = cute.arch.warp_idx() + rows: cutlass.Constexpr = self.rows + outs: cutlass.Constexpr = self.outs + threads: cutlass.Constexpr = self.threads + k_tiles: cutlass.Constexpr = self.k_tiles + warps: cutlass.Constexpr = self.warps + n0 = bidx * outs + my_col = n0 + tidx + + # (row, (VEC, (threads, k_tiles))): thread t's vectors of a row, k-tile by k-tile. + w_vec = cute.logical_divide(cute.logical_divide(mW, (None, VEC)), (None, (None, threads))) + a_vec = cute.logical_divide(cute.logical_divide(mA, (None, VEC)), (None, (None, threads))) + + # 1. Static operands, before the wait. + w_reg = cute.make_rmem_tensor( + cute.make_layout((outs, k_tiles, VEC), stride=(k_tiles * VEC, VEC, 1)), cutlass.BFloat16 + ) + for j in cutlass.range_constexpr(outs): + w_row = w_vec[n0 + j, (None, (tidx, None))] + for kt in cutlass.range_constexpr(k_tiles): + cute.copy(load_w, w_row[None, kt], w_reg[j, kt, None]) + tap = cute.make_rmem_tensor(cute.make_layout((WIDTH,), stride=(1,)), cutlass.BFloat16) + win = cute.make_rmem_tensor( + cute.make_layout((rows, WIDTH - 1), stride=(WIDTH - 1, 1)), cutlass.BFloat16 + ) + if tidx < outs: + if my_col < self.conv_dim: + for j4 in cutlass.range_constexpr(WIDTH): + tap[j4] = mConvW[my_col, j4] + for r in cutlass.range_constexpr(rows): + slot = mSlot[r] + if slot != self.pad_slot: + for j3 in cutlass.range_constexpr(WIDTH - 1): + win[r, j3] = mState[slot, my_col, j3] + + if const_expr(self.pdl): + n_rows = cute.size(mW, mode=[0]) + k_len = cute.size(mW, mode=[1]) + if tidx < L2_AHEAD: + tgt = bidx + (tidx + 1) * L2_DIST + if (tgt + 1) * outs <= n_rows: + _bulk_prefetch_l2(mW.iterator + tgt * outs * k_len, outs * k_len * 2) + + # 2. The producer's output only after the wait. + if const_expr(self.pdl): + cute.arch.griddepcontrol_launch_dependents() + cute.arch.griddepcontrol_wait() + + # 3. Products and reductions. + acc = cute.make_rmem_tensor( + cute.make_layout((rows, outs), stride=(outs, 1)), cutlass.Float32 + ) + acc.fill(0.0) + a_reg = cute.make_rmem_tensor( + cute.make_layout((k_tiles, VEC), stride=(VEC, 1)), cutlass.BFloat16 + ) + for r in cutlass.range_constexpr(rows): + a_row = a_vec[r, (None, (tidx, None))] + for kt in cutlass.range_constexpr(k_tiles): + cute.copy(load_a, a_row[None, kt], a_reg[kt, None]) + for kt in cutlass.range_constexpr(k_tiles): + for v in cutlass.range_constexpr(VEC): + av = a_reg[kt, v].to(cutlass.Float32) + for j in cutlass.range_constexpr(outs): + acc[r, j] = acc[r, j] + av * w_reg[j, kt, v].to(cutlass.Float32) + for r in cutlass.range_constexpr(rows): + for j in cutlass.range_constexpr(outs): + acc[r, j] = cute.arch.warp_reduction_sum(acc[r, j]) + smem = cutlass.utils.SmemAllocator() + part = smem.allocate_tensor( + cutlass.Float32, + cute.make_layout((rows, outs, warps), stride=(outs * warps, warps, 1)), + byte_alignment=16, + ) + with cute.arch.elect_one(): + for r in cutlass.range_constexpr(rows): + for j in cutlass.range_constexpr(outs): + part[r, j, warp] = acc[r, j] + cute.arch.sync_threads() + + # 4. Thread j < outs finishes column n0 + j. + if tidx < outs: + for r in cutlass.range_constexpr(rows): + total = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(warps): + total = total + part[r, tidx, w] + proj = total.to(cutlass.BFloat16) + res = proj + if my_col < self.conv_dim: + slot = mSlot[r] + if slot != self.pad_slot: + c = ( + tap[0].to(cutlass.Float32) * win[r, 0].to(cutlass.Float32) + + tap[1].to(cutlass.Float32) * win[r, 1].to(cutlass.Float32) + + tap[2].to(cutlass.Float32) * win[r, 2].to(cutlass.Float32) + + tap[3].to(cutlass.Float32) * proj.to(cutlass.Float32) + ) + res = (c / (1.0 + cute.math.exp(-c))).to(cutlass.BFloat16) + mState[slot, my_col, 0] = win[r, 1] + mState[slot, my_col, 1] = win[r, 2] + mState[slot, my_col, 2] = proj + mOut[r, my_col] = res + + +def _dl(t: torch.Tensor, align: int): + return from_dlpack(t, assumed_align=align) + + +@functools.lru_cache(maxsize=None) +def _compiled(rows, n, k, conv_dim, pad_slot, pdl, state_shape, state_stride): + threads, outs = TILING[k] + kern = _InProjConvKernel( + rows=rows, k=k, threads=threads, outs=outs, conv_dim=conv_dim, pad_slot=pad_slot, pdl=pdl + ) + dev = "cuda" + a = torch.empty(rows, k, dtype=torch.bfloat16, device=dev) + w = torch.empty(n, k, dtype=torch.bfloat16, device=dev) + out = torch.empty(rows, n, dtype=torch.bfloat16, device=dev) + state = torch.empty_strided(state_shape, state_stride, dtype=torch.bfloat16, device=dev) + convw = torch.empty(conv_dim, WIDTH, dtype=torch.bfloat16, device=dev) + slot = torch.empty(rows, dtype=torch.int32, device=dev) + stream = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream) + return cute_ext.compile( + kern, + _dl(a, 32), + _dl(w, 32), + _dl(out, 32), + _dl(state, 2), + _dl(convw, 8), + _dl(slot, 4), + stream, + options=_COMPILE_OPTIONS, + ) + + +def launch(x, mat_b, conv_state, conv_weight, conv_state_indices, conv_dim, pad_slot_id, pdl): + """The launch behind `GdnInProjConv1dPdl.__call__`. Returns the fresh `[M, N]` output.""" + rows, k = x.shape + n = mat_b.shape[1] + dev = x.device + pdl = pdl and pdl_enabled(dev.index if dev.index is not None else torch.cuda.current_device()) + out = torch.empty(rows, n, dtype=x.dtype, device=dev) + fn = _compiled( + rows, + n, + k, + conv_dim, + int(pad_slot_id), + bool(pdl), + tuple(conv_state.shape), + tuple(conv_state.stride()), + ) + stream = cuda_drv.CUstream(torch.cuda.current_stream(dev).cuda_stream) + fn( + _dl(x, 32), + _dl(mat_b.t(), 32), + _dl(out, 32), + _dl(conv_state, 2), + _dl(conv_weight, 8), + _dl(conv_state_indices, 4), + stream, + ) + return out + + +def _cell(why: str, rows: int, slots, pdl: bool = True, pool: int = 461) -> Cell: + return Cell(why=why, spec=dict(rows=rows, slots=slots, pdl=pdl, pool=pool)) + + +class GdnInProjConv1dPdl(OpWrapper): + """`out = x @ mat_b`, then for `c < conv_dim` and a row whose slot is not + `pad_slot_id` the causal conv step on the projected value, one launch. + + `x` is the dense row-major `[M, K]` activation (`M <= MAX_ROWS`), `mat_b` the + dense column-major `[K, N]` view of the stored `[N, K]` projection (its + `.t()`). Row `r`, slot `s = conv_state_indices[r]`, column `c < conv_dim`: + `p = round_bf16(sum_k x[r, k] * W[c, k])` (fp32 products and sums), + `out[r, c] = round_bf16(silu(sum_j conv_weight[c, j] * [conv_state[s, c, :], p][j]))` + in fp32, and `conv_state[s, c, :] = [conv_state[s, c, 1:], p]`. Every other + column, and every column of a pad row, is `p`; a pad row's slot is not read or + written. Returns the fresh `[M, N]` output; `conv_state` is advanced in place. + + Equivalent to `low_m_bf16_mm` followed by `causal_conv1d_update(x=out[:, :conv_dim], + conv_state, conv_weight, bias=None, silu_activation=True, cache_seqlens=None, + conv_state_indices, pad_slot_id)` up to the GEMV's summation order. + + Fusion boundary: the projection, the conv step, SiLU and the window's advance. + The split into q|k|v, z, a, b and everything after are the caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: rows, slots (tuple per row; -1 is the pad slot), pdl, pool (slots in the pool). + CELLS: tuple[Cell, ...] = ( + _cell("decode 1 row, slot 37 of the 461-slot pool", 1, (37,)), + _cell("decode 1 row, the last slot of the pool", 1, (460,)), + _cell("decode 2 rows, two slots", 2, (5, 211)), + _cell("decode 2 rows, one a pad row (-1)", 2, (17, -1)), + _cell("decode 1 row, PDL off", 1, (99,), pdl=False), + _cell("decode 1 row, a 129-slot pool", 1, (128,), pool=129), + ) + + note = """ + Metadata consumed: the decode rows' conv slots (`conv_state_indices`, int32 + `[M]`). Preconditions: bf16 `x` dense `[M, K]`, `1 <= M <= MAX_ROWS`; `mat_b` + the dense column-major `[K, N]` view of a `[N, K]` weight with K in `TILING` + and `N` a multiple of the tile's outputs; `conv_state` `[slots, conv_dim, 3]` + bf16 (any strides; they are compiled in), `conv_weight` `[conv_dim, 4]` bf16 + contiguous, `conv_dim <= N`, no bias, SiLU on, `cache_seqlens` None (the + window shifts; not a ring). + + PDL: the weight rows, the conv taps, the slots and the pool's windows are read + before `griddepcontrol.wait` (the pool row is this layer's own, written by this + kernel one step ago; nothing between writes it); the activation after it. The + launch triggers its dependents right before its own wait. + + Compilation: CuTe DSL, JIT on the first call per (M, N, K, conv_dim, + pad_slot_id, pdl, pool shape and strides); that call must come before + CUDA-graph capture -- the engine's warmup forward does it. + + Numerics: the projection is fp32 products and sums rounded once to bf16 (the + value stored for non-conv columns and pushed into the window); the conv is + fp32 over the bf16 window, taps and projection, rounded once. Only the + GEMV's summation order differs from `low_m_bf16_mm` + `causal_conv1d_update`. + + Hard limits + M rows per launch boundary MAX_ROWS=2 driven at 1, 2 (cells) + K a TILING key 2560 the checkpoint's hidden size + pool any slot count and strides 461 and 129 driven + """ + + def __call__( + self, + x: torch.Tensor, + mat_b: torch.Tensor, + conv_state: torch.Tensor, + conv_weight: torch.Tensor, + conv_state_indices: torch.Tensor, + conv_dim: int, + pad_slot_id: int, + pdl: bool = True, + ) -> torch.Tensor: + return launch( + x, mat_b, conv_state, conv_weight, conv_state_indices, conv_dim, pad_slot_id, pdl + ) + + def reference( + self, + x: torch.Tensor, + mat_b: torch.Tensor, + conv_state: torch.Tensor, + conv_weight: torch.Tensor, + conv_state_indices: torch.Tensor, + conv_dim: int, + pad_slot_id: int, + pdl: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + """`(out, conv_state)` as the call leaves them, in fp32 then rounded; mutates nothing.""" + with true_fp32_matmul(): + proj = (x.float() @ mat_b.float()).to(x.dtype) + out = proj.clone() + state = conv_state.clone() + for r, s in enumerate(conv_state_indices.tolist()): + if s == pad_slot_id: + continue + p = proj[r, :conv_dim].float() + window = torch.cat([conv_state[s].float(), p[:, None]], dim=1) + c = (window * conv_weight.float()).sum(1) + out[r, :conv_dim] = (c * torch.sigmoid(c)).to(x.dtype) + state[s] = window[:, 1:].to(conv_state.dtype) + return out, state + + def is_valid( + self, + x: torch.Tensor, + mat_b: torch.Tensor, + conv_state: torch.Tensor, + conv_weight: torch.Tensor, + conv_state_indices: torch.Tensor, + conv_dim: int, + pad_slot_id: int, + pdl: bool = True, + ) -> None: + assert x.dtype == mat_b.dtype == conv_state.dtype == conv_weight.dtype == torch.bfloat16 + assert x.dim() == 2 and x.is_contiguous() and 1 <= x.shape[0] <= MAX_ROWS, tuple(x.shape) + assert mat_b.shape[0] == x.shape[1] and x.shape[1] in TILING, ( + tuple(mat_b.shape), + tuple(x.shape), + ) + assert mat_b.stride(0) == 1 and mat_b.stride(1) == mat_b.shape[0], ( + "mat_b must be dense column-major [K, N] (e.g. weight.t())" + ) + assert mat_b.shape[1] % TILING[x.shape[1]][1] == 0, ( + "N must be a multiple of the tile's outputs" + ) + assert conv_state.dim() == 3 and tuple(conv_state.shape[1:]) == (conv_dim, WIDTH - 1) + assert tuple(conv_weight.shape) == (conv_dim, WIDTH) and conv_weight.is_contiguous() + assert conv_dim <= mat_b.shape[1] + assert conv_state_indices.dtype == torch.int32 and tuple(conv_state_indices.shape) == ( + x.shape[0], + ) + + +gdn_in_proj_conv1d_pdl = GdnInProjConv1dPdl() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_prefill_conv_post.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_prefill_conv_post.py new file mode 100644 index 000000000000..06b7863fc652 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_prefill_conv_post.py @@ -0,0 +1,1285 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Gated-DeltaNet prefill preparation in one pass: the causal depthwise conv1d +of the context rows, read token-major straight off the input projection, with +the SiLU, the q/k split and L2 norm, the v copy and the decay/beta gates fused +in, and each sequence's conv window published to the paged pool. + +Tier 6 (CuTe DSL written for this target) for the conv pass, Tier 5 (Triton) +for the gates, the window publish and the fallback conv pass. The three-call +composition it replaces -- `extract_transpose_prefill_slice` (token-major -> +channel-major copy), `causal_conv1d_fwd` (channel-major conv, in place) and +`fused_gdn_post_conv` (split + norm + gates) -- moves the `[T, 10240]` +activation through HBM three times; no ``torch.ops.trtllm.*`` op, in-tree +callable or installed FlashInfer API fuses the conv with the norm, and the +channel-last `causal_conv1d_fwd` kernel alone is slower than this whole pass. +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +import triton +import triton.language as tl +from cutlass._mlir.dialects import llvm +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.cute_dsl_kernels.blackwell.utils import make_ptr + +#: This checkpoint's linear-attention geometry at tp1. +NUM_K_HEADS, HEAD_K_DIM, NUM_V_HEADS, HEAD_V_DIM = 16, 128, 48, 128 +QKV_DIM = 2 * NUM_K_HEADS * HEAD_K_DIM + NUM_V_HEADS * HEAD_V_DIM # 10240 +#: The target's one GDN input-projection output: [q|k|v | z | a | b]. +IN_PROJ_WIDTH = QKV_DIM + NUM_V_HEADS * HEAD_V_DIM + 2 * NUM_V_HEADS # 16480 +CONV_WIDTH = 4 +#: Largest conv-pool slot count the cache manager was seen to allocate for this +#: checkpoint (the `causal_conv1d_fwd` entry's ENGINE_SLOTS). +ENGINE_SLOTS = 461 +A_LOG_RANGE = (-5.5625, 3.5625) +DT_BIAS_RANGE = (-8.9375, 19.25) +SOFTPLUS_THRESHOLD = 20.0 + +#: Launch shape of the main kernel, chosen by micro-bench at 2048, 8192 and +#: 16384 tokens (see the entry note). The register cap is what makes it fast: +#: uncapped, the compiler takes 140+ registers for the pool-reading branch and +#: the occupancy that leaves cannot hide the load latency; capped at 72 that +#: branch spills, and it runs only on blocks holding a sequence's first rows. +_BLOCK_T = 32 +_NUM_WARPS = 4 +_MAXNREG = 72 +#: Channels per program of the window publish. +_PUBLISH_BLOCK_C = 512 + + +@triton.jit(do_not_specialize=["num_tokens", "num_seqs", "search_steps"]) +def _gdn_prefill_conv_post_kernel( + x_ptr, + stride_xt, + w_ptr, + stride_wc, + stride_wk, + state_ptr, + stride_ss, + stride_sc, + stride_sw, + qsl_ptr, + slot_ptr, + init_ptr, + a_ptr, + stride_at, + stride_ah, + b_ptr, + stride_bt, + stride_bh, + A_log_ptr, + dt_bias_ptr, + q_ptr, + k_ptr, + v_ptr, + g_ptr, + beta_ptr, + num_tokens, + num_seqs, + search_steps, + pad_slot_id, + l2_norm_eps, + softplus_threshold, + IDX: tl.constexpr, + NUM_K_HEADS: tl.constexpr, + NUM_V_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_T: tl.constexpr, +): + """One program: BLOCK_T tokens x one 128-channel head of the conv output.""" + head = tl.program_id(0) + t = tl.program_id(1) * BLOCK_T + tl.arange(0, BLOCK_T) + t_mask = t < num_tokens + d = tl.arange(0, HEAD_DIM) + c = head * HEAD_DIM + d + rows2 = (tl.program_id(1) * BLOCK_T + tl.arange(0, BLOCK_T))[:, None] + rmask2 = rows2 < num_tokens + + # Sequence of each token: the largest s with qsl[s] <= t (a zero-length + # sequence shares its start with the next one and is never chosen over it). + lo = tl.zeros([BLOCK_T], dtype=tl.int32) + hi = tl.zeros([BLOCK_T], dtype=tl.int32) + num_seqs + for _ in range(search_steps): + mid = (lo + hi) // 2 + right = tl.load(qsl_ptr + mid) <= t + lo = tl.where(right, mid, lo) + hi = tl.where(right, hi, mid) + s0 = tl.load(qsl_ptr + lo) + + # Width-4 causal conv, taps accumulated in fp32 in tap order. A tap before + # the sequence start reads the carried window column W-1-(s0-pos), or 0. + # Only a block holding one of a sequence's first three tokens needs that; + # every other block takes the branch whose taps are plain row loads. + acc = tl.zeros([BLOCK_T, HEAD_DIM], dtype=tl.float32) + if tl.min(tl.where(t_mask, t - s0, 3), axis=0) < 3: + slot = tl.load(slot_ptr + lo, mask=t_mask, other=pad_slot_id) + carry = t_mask & (tl.load(init_ptr + lo, mask=t_mask, other=0) != 0) & (slot != pad_slot_id) + slot64 = slot.to(tl.int64) + for kk in tl.static_range(4): + pos = t - 3 + kk + in_seq = pos >= s0 + xv = tl.load( + x_ptr + pos[:, None].to(IDX) * stride_xt + c[None, :], + mask=(t_mask & in_seq)[:, None], + other=0.0, + ) + s_row = slot64 * stride_ss + (3 + pos - s0) * stride_sw + sv = tl.load( + state_ptr + s_row[:, None] + c[None, :] * stride_sc, + mask=(carry & (~in_seq))[:, None], + other=0.0, + ) + wk = tl.load(w_ptr + c * stride_wc + kk * stride_wk).to(tl.float32) + acc += tl.where(in_seq[:, None], xv, sv).to(tl.float32) * wk[None, :] + else: + for kk in tl.static_range(4): + xv = tl.load( + x_ptr + (rows2 - (3 - kk)).to(IDX) * stride_xt + c[None, :], mask=rmask2, other=0.0 + ) + wk = tl.load(w_ptr + c * stride_wc + kk * stride_wk).to(tl.float32) + acc += xv.to(tl.float32) * wk[None, :] + act = acc / (1.0 + tl.exp(-acc)) + # The conv result is rounded to the activation dtype once, as the + # unfused conv stores it, before anything reads it back. + y = act.to(q_ptr.dtype.element_ty) + + ri2 = rows2.to(IDX) + if head < 2 * NUM_K_HEADS: + yf = y.to(tl.float32) + inv = 1.0 / tl.sqrt(tl.sum(yf * yf, axis=1) + l2_norm_eps) + normed = (yf * inv[:, None]).to(q_ptr.dtype.element_ty) + o_off = ri2 * (NUM_K_HEADS * HEAD_DIM) + (head % NUM_K_HEADS) * HEAD_DIM + d[None, :] + if head < NUM_K_HEADS: + tl.store(q_ptr + o_off, normed, mask=rmask2) + else: + tl.store(k_ptr + o_off, normed, mask=rmask2) + else: + vh = head - 2 * NUM_K_HEADS + o_off = ri2 * (NUM_V_HEADS * HEAD_DIM) + vh * HEAD_DIM + d[None, :] + tl.store(v_ptr + o_off, y, mask=rmask2) + t64 = t.to(IDX) + a_val = tl.load(a_ptr + t64 * stride_at + vh * stride_ah, mask=t_mask, other=0.0) + b_val = tl.load(b_ptr + t64 * stride_bt + vh * stride_bh, mask=t_mask, other=0.0) + A_log = tl.load(A_log_ptr + vh).to(tl.float32) + dt_bias = tl.load(dt_bias_ptr + vh).to(tl.float32) + gate_input = a_val.to(tl.float32) + dt_bias + softplus = tl.where( + gate_input <= softplus_threshold, tl.log(1.0 + tl.exp(gate_input)), gate_input + ) + g_off = t64 * NUM_V_HEADS + vh + tl.store(g_ptr + g_off, -tl.exp(A_log) * softplus, mask=t_mask) + tl.store(beta_ptr + g_off, tl.sigmoid(b_val.to(tl.float32)), mask=t_mask) + + +@triton.jit +def _gdn_prefill_conv_publish_kernel( + x_ptr, + stride_xt, + state_ptr, + stride_ss, + stride_sc, + stride_sw, + qsl_ptr, + slot_ptr, + init_ptr, + conv_dim, + pad_slot_id, + BLOCK_C: tl.constexpr, +): + """One program: one sequence x BLOCK_C channels of its new 3-column window. + + New column j holds position `e - 3 + j`: the raw input when it lies inside + the sequence, else the carried column `L + j` (or 0 without a carry). All + three are loaded before any is stored, and each channel's three columns + live in one thread, so a short sequence's shift reads the old window.""" + n = tl.program_id(0) + c = tl.program_id(1) * BLOCK_C + tl.arange(0, BLOCK_C) + c_mask = c < conv_dim + slot = tl.load(slot_ptr + n) + if slot != pad_slot_id: + s0 = tl.load(qsl_ptr + n) + e = tl.load(qsl_ptr + n + 1) + carry = tl.load(init_ptr + n) != 0 + base = state_ptr + slot.to(tl.int64) * stride_ss + c.to(tl.int64) * stride_sc + p0 = e - 3 + p1 = e - 2 + p2 = e - 1 + x0 = tl.load(x_ptr + p0.to(tl.int64) * stride_xt + c, mask=c_mask & (p0 >= s0), other=0.0) + x1 = tl.load(x_ptr + p1.to(tl.int64) * stride_xt + c, mask=c_mask & (p1 >= s0), other=0.0) + x2 = tl.load(x_ptr + p2.to(tl.int64) * stride_xt + c, mask=c_mask & (p2 >= s0), other=0.0) + s_0 = tl.load(base + (3 + p0 - s0) * stride_sw, mask=c_mask & (p0 < s0) & carry, other=0.0) + s_1 = tl.load(base + (3 + p1 - s0) * stride_sw, mask=c_mask & (p1 < s0) & carry, other=0.0) + s_2 = tl.load(base + (3 + p2 - s0) * stride_sw, mask=c_mask & (p2 < s0) & carry, other=0.0) + n0 = tl.where(p0 >= s0, x0, s_0) + n1 = tl.where(p1 >= s0, x1, s_1) + n2 = tl.where(p2 >= s0, x2, s_2) + tl.store(base, n0, mask=c_mask) + tl.store(base + stride_sw, n1, mask=c_mask) + tl.store(base + 2 * stride_sw, n2, mask=c_mask) + + +# -------------------------------------------------------------------------------------- +# The conv pass in CuTe DSL. A CTA is one 128-channel head x `_TT` tokens, 128 threads: +# it first copies its `_TT + 3`-row tile (3 history rows) into shared memory with +# 16-byte `cp.async` (rows outside `[0, T)` zero-filled), so the bytes in flight cost no +# registers; then a thread (4 channels x `_TPT` consecutive tokens) slides the width-4 +# window over shared-memory rows, every element converted to fp32 once. The conv runs +# as `fma.rn.f32x2` on channel pairs (first tap from zero, then in tap order: per element +# the Triton pass's `fma.rn.f32` chain); the SiLU is `acc * rcp(1 + 2^(-acc * log2e))` +# with `ex2.approx.ftz` and `rcp.approx.ftz`, rounded in pairs by `cvt.rn.bf16x2.f32`. +# -------------------------------------------------------------------------------------- + +_VEC = 4 +_TPT = 16 +_RG = 4 +_TT = _TPT * _RG +_TILE_ROWS = _TT + 3 +_THREADS = 128 +_CHUNKS = _TILE_ROWS * 16 + + +def _iv(x, loc, ip): + return x.ir_value(loc=loc, ip=ip) + + +def _asm1(ret, args, text, constraints, loc, ip, side_effects=False): + return llvm.inline_asm( + ret, + [_iv(a, loc, ip) for a in args], + text, + constraints, + has_side_effects=side_effects, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _asm2(ty, args, text, constraints, loc, ip, side_effects=False): + res = _asm1( + llvm.StructType.get_literal([ty, ty]), args, text, constraints, loc, ip, side_effects + ) + return ( + llvm.extractvalue(ty, res, [0], loc=loc, ip=ip), + llvm.extractvalue(ty, res, [1], loc=loc, ip=ip), + ) + + +@dsl_user_op +def _ffma2(a0, a1, b0, b1, c0, c1, *, loc=None, ip=None): + """`fma.rn.f32x2`: two fp32 FMAs, each rounded once.""" + r = _asm2( + T.f32(), + [cutlass.Float32(v) for v in (a0, a1, b0, b1, c0, c1)], + "{ .reg .b64 a, b, c, d; mov.b64 a, {$2, $3}; mov.b64 b, {$4, $5}; " + "mov.b64 c, {$6, $7}; fma.rn.f32x2 d, a, b, c; mov.b64 {$0, $1}, d; }", + "=f,=f,f,f,f,f,f,f", + loc, + ip, + ) + return cutlass.Float32(r[0]), cutlass.Float32(r[1]) + + +@dsl_user_op +def _fmul2(a0, a1, b0, b1, *, loc=None, ip=None): + """`mul.rn.f32x2`.""" + r = _asm2( + T.f32(), + [cutlass.Float32(v) for v in (a0, a1, b0, b1)], + "{ .reg .b64 a, b, d; mov.b64 a, {$2, $3}; mov.b64 b, {$4, $5}; " + "mul.rn.f32x2 d, a, b; mov.b64 {$0, $1}, d; }", + "=f,=f,f,f,f,f", + loc, + ip, + ) + return cutlass.Float32(r[0]), cutlass.Float32(r[1]) + + +@dsl_user_op +def _silu2(a0, a1, *, loc=None, ip=None): + """`acc * rcp(1 + 2^(acc * -log2e))` for two values: `mul.rn.f32x2`, `ex2.approx.ftz`, + `add.rn.f32x2`, `rcp.approx.ftz`, `mul.rn`. For `1 + e` in [1, 2^126] this equals the + Triton pass's `div.full.f32`; past it (acc < -87.3) it gives -0 where that gives a + value below 1e-36.""" + r = _asm2( + T.f32(), + [cutlass.Float32(a0), cutlass.Float32(a1)], + "{ .reg .b64 a, k, t, o, d; .reg .f32 c, t0, t1, d0, d1; " + "mov.f32 c, 0fBFB8AA3B; mov.b64 k, {c, c}; mov.f32 c, 0f3F800000; mov.b64 o, {c, c}; " + "mov.b64 a, {$2, $3}; mul.rn.f32x2 t, a, k; mov.b64 {t0, t1}, t; " + "ex2.approx.ftz.f32 t0, t0; ex2.approx.ftz.f32 t1, t1; mov.b64 t, {t0, t1}; " + "add.rn.f32x2 d, t, o; mov.b64 {d0, d1}, d; " + "rcp.approx.ftz.f32 d0, d0; rcp.approx.ftz.f32 d1, d1; " + "mul.rn.f32 $0, $2, d0; mul.rn.f32 $1, $3, d1; }", + "=f,=f,f,f", + loc, + ip, + ) + return cutlass.Float32(r[0]), cutlass.Float32(r[1]) + + +@dsl_user_op +def _fma_rn(a, b, c, *, loc=None, ip=None): + return cutlass.Float32( + _asm1( + T.f32(), + [cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _inv_norm(ss, eps, *, loc=None, ip=None): + """1 / sqrt(ss + eps): `add.f32`, `sqrt.approx.ftz.f32`, `div.full.f32` (the Triton pass's).""" + return cutlass.Float32( + _asm1( + T.f32(), + [cutlass.Float32(ss), cutlass.Float32(eps)], + "{ .reg .f32 t; add.f32 t, $1, $2; sqrt.approx.ftz.f32 t, t; " + "div.full.f32 $0, 0f3F800000, t; }", + "=f,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _bf16x2(lo, hi, *, loc=None, ip=None): + """`cvt.rn.bf16x2.f32`: `lo` in the low half (the lower address).""" + return cutlass.Int32( + _asm1( + T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "cvt.rn.bf16x2.f32 $0, $1, $2;", + "=r,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _unpack2(w, *, loc=None, ip=None): + """A bf16x2 word as `(low, high)` fp32.""" + r = _asm2( + T.f32(), + [cutlass.Int32(w)], + "{ .reg .b32 l, h; shl.b32 l, $2, 16; and.b32 h, $2, 0xFFFF0000; " + "mov.b32 $0, l; mov.b32 $1, h; }", + "=f,=f,r", + loc, + ip, + ) + return cutlass.Float32(r[0]), cutlass.Float32(r[1]) + + +@dsl_user_op +def _st_v2(addr, a, b, *, loc=None, ip=None): + _asm1( + T.i32(), + [cutlass.Int64(addr), cutlass.Int32(a), cutlass.Int32(b)], + "{ st.global.v2.b32 [$1], {$2, $3}; mov.u32 $0, 0; }", + "=r,l,r,r", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _cp_async16(smem_addr, gmem_addr, src_bytes, *, loc=None, ip=None): + """16 bytes global -> shared, the rest of the 16 zero-filled past `src_bytes`.""" + _asm1( + T.i32(), + [cutlass.Int32(smem_addr), cutlass.Int64(gmem_addr), cutlass.Int32(src_bytes)], + "{ cp.async.cg.shared.global [$1], [$2], 16, $3; mov.u32 $0, 0; }", + "=r,r,l,r", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _cp_async_wait_all(*, loc=None, ip=None): + _asm1(T.i32(), [], "{ cp.async.wait_all; mov.u32 $0, 0; }", "=r", loc, ip, side_effects=True) + + +@dsl_user_op +def _lds_v2(addr, *, loc=None, ip=None): + r = _asm2( + T.i32(), + [cutlass.Int32(addr)], + "ld.shared.v2.b32 {$0, $1}, [$2];", + "=r,=r,r", + loc, + ip, + side_effects=True, + ) + return cutlass.Int32(r[0]), cutlass.Int32(r[1]) + + +class _ConvPassKernel: + """q / k / v of the conv pass; the gates and the publish are separate launches.""" + + def __init__(self, num_k_heads: int, num_v_heads: int, head_dim: int, x_stride: int): + assert head_dim == _VEC * 32 + self.hk, self.hv, self.d = num_k_heads, num_v_heads, head_dim + self.heads = 2 * num_k_heads + num_v_heads + self.qkv = self.heads * head_dim + self.sx = x_stride + + @cute.jit + def __call__( + self, + x_ptr: cute.Pointer, + w_ptr: cute.Pointer, + st_ptr: cute.Pointer, + qsl_ptr: cute.Pointer, + slot_ptr: cute.Pointer, + init_ptr: cute.Pointer, + q_t: cute.Tensor, + k_ptr: cute.Pointer, + v_ptr: cute.Pointer, + num_tokens: cutlass.Int32, + num_seqs: cutlass.Int32, + search_steps: cutlass.Int32, + pad_slot_id: cutlass.Int32, + eps: cutlass.Float32, + w_s0: cutlass.Int32, + w_s1: cutlass.Int32, + st_s0: cutlass.Int32, + st_s1: cutlass.Int32, + st_s2: cutlass.Int32, + stream: cuda.CUstream, + ): + weight = cute.make_tensor(w_ptr, cute.make_layout((self.qkv, 4), stride=(w_s0, w_s1))) + # The pool's slot extent is never read: only `cache_indices` slots are addressed. + pool = cute.make_tensor( + st_ptr, cute.make_layout((1 << 20, self.qkv, 3), stride=(st_s0, st_s1, st_s2)) + ) + qsl = cute.make_tensor(qsl_ptr, cute.make_layout((num_seqs + 1,), stride=(1,))) + slots = cute.make_tensor(slot_ptr, cute.make_layout((num_seqs,), stride=(1,))) + init = cute.make_tensor(init_ptr, cute.make_layout((num_seqs,), stride=(1,))) + blocks = (num_tokens + _TT - 1) // _TT + # `q` arrives as a tensor: the TVM-FFI entry takes its stream from a GPU tensor. + self.kernel( + x_ptr.toint(), + weight, + pool, + qsl, + slots, + init, + q_t.iterator.toint(), + k_ptr.toint(), + v_ptr.toint(), + num_tokens, + num_seqs, + search_steps, + pad_slot_id, + eps, + ).launch( + grid=[self.heads, blocks, 1], + block=[_THREADS, 1, 1], + smem=_TILE_ROWS * 128 * 2 + 128, + stream=stream, + ) + + @cute.kernel + def kernel( + self, + x_base: cutlass.Int64, + mW: cute.Tensor, + mState: cute.Tensor, + mQsl: cute.Tensor, + mSlot: cute.Tensor, + mInit: cute.Tensor, + q_base: cutlass.Int64, + k_base: cutlass.Int64, + v_base: cutlass.Int64, + num_tokens: cutlass.Int32, + num_seqs: cutlass.Int32, + search_steps: cutlass.Int32, + pad_slot_id: cutlass.Int32, + eps: cutlass.Float32, + ): + tidx, _, _ = cute.arch.thread_idx() + head, tb, _ = cute.arch.block_idx() + lane = tidx % 32 + rg = tidx // 32 + t0 = tb * _TT + + smem = cutlass.utils.SmemAllocator() + tile = smem.allocate_tensor( + cutlass.BFloat16, cute.make_layout((_TILE_ROWS * 128,), stride=(1,)), byte_alignment=16 + ) + tile_base = cutlass.Int32(tile.iterator.toint()) + + # 1. The tile: rows t0-3 .. t0+TT-1 of this head's 128 channels, zero-filled outside [0, T). + for it in cutlass.range_constexpr((_CHUNKS + _THREADS - 1) // _THREADS): + chunk = tidx + it * _THREADS + if chunk < _CHUNKS: + r = chunk // 16 + cc = chunk - r * 16 + row = t0 - 3 + r + valid = cutlass.Int32(0) + if row >= 0: + if row < num_tokens: + valid = cutlass.Int32(16) + rowc = cutlass.min(cutlass.max(row, cutlass.Int32(0)), num_tokens - 1) + _cp_async16( + tile_base + (r * 128 + cc * 8) * 2, + x_base + (cutlass.Int64(rowc) * self.sx + head * 128 + cc * 8) * 2, + valid, + ) + + # 2. While the copies land: the weights, the first token's sequence, the output row. + c0 = head * 128 + lane * _VEC + w = cute.make_rmem_tensor(cute.make_layout((_VEC, 4), stride=(4, 1)), cutlass.Float32) + for i in cutlass.range_constexpr(_VEC * 4): + w[i // 4, i % 4] = mW[c0 + i // 4, i % 4].to(cutlass.Float32) + tf = t0 + rg * _TPT + lo = cutlass.Int32(0) + hi = num_seqs + for _ in cutlass.range(search_steps, unroll=1): + mid = (lo + hi) // 2 + if mQsl[mid] <= tf: + lo = mid + else: + hi = mid + s0f = mQsl[lo] + ef = mQsl[lo + 1] + # One uniform branch a thread: its tokens and their windows inside one sequence. + fast = cutlass.Int32(0) + if tf - s0f >= 3: + if tf + _TPT <= ef: + if tf + _TPT <= num_tokens: + fast = cutlass.Int32(1) + is_qk = head < 2 * self.hk + o_base = v_base + cutlass.Int64((head - 2 * self.hk) * 128 + lane * _VEC) * 2 + o_stride = cutlass.Int64(self.hv * 128 * 2) + if head < self.hk: + o_base = q_base + cutlass.Int64(head * 128 + lane * _VEC) * 2 + o_stride = cutlass.Int64(self.hk * 128 * 2) + else: + if is_qk: + o_base = k_base + cutlass.Int64((head - self.hk) * 128 + lane * _VEC) * 2 + o_stride = cutlass.Int64(self.hk * 128 * 2) + + _cp_async_wait_all() + cute.arch.sync_threads() + + # 3. The window over tile rows rg*TPT .. (tile row r is token t0 - 3 + r). + win = cute.make_rmem_tensor(cute.make_layout((4, _VEC), stride=(_VEC, 1)), cutlass.Float32) + for i in cutlass.range_constexpr(3): + lw, hw = _lds_v2(tile_base + ((rg * _TPT + i) * 128 + lane * _VEC) * 2) + win[i, 0], win[i, 1] = _unpack2(lw) + win[i, 2], win[i, 3] = _unpack2(hw) + acc = cute.make_rmem_tensor(cute.make_layout((_VEC,), stride=(1,)), cutlass.Float32) + for j in cutlass.range_constexpr(_TPT): + t = tf + j + lw, hw = _lds_v2(tile_base + ((rg * _TPT + j + 3) * 128 + lane * _VEC) * 2) + win[(j + 3) % 4, 0], win[(j + 3) % 4, 1] = _unpack2(lw) + win[(j + 3) % 4, 2], win[(j + 3) % 4, 3] = _unpack2(hw) + if fast == 1: + for p in cutlass.range_constexpr(_VEC // 2): + a0, a1 = _ffma2( + win[j % 4, 2 * p], + win[j % 4, 2 * p + 1], + w[2 * p, 0], + w[2 * p + 1, 0], + cutlass.Float32(0.0), + cutlass.Float32(0.0), + ) + for kk in cutlass.range_constexpr(1, 4): + a0, a1 = _ffma2( + win[(j + kk) % 4, 2 * p], + win[(j + kk) % 4, 2 * p + 1], + w[2 * p, kk], + w[2 * p + 1, kk], + a0, + a1, + ) + acc[2 * p] = a0 + acc[2 * p + 1] = a1 + else: + # The token's sequence; a tap before its start reads the carried + # column W-1-(s0-pos), or 0 without a carry. + sq = lo + s0 = s0f + if t >= ef: + lo2 = cutlass.Int32(0) + hi2 = num_seqs + for _ in cutlass.range(search_steps, unroll=1): + mid = (lo2 + hi2) // 2 + if mQsl[mid] <= t: + lo2 = mid + else: + hi2 = mid + sq = lo2 + s0 = mQsl[lo2] + slot = mSlot[sq] + init = cutlass.Int32(mInit[sq]) + for v in cutlass.range_constexpr(_VEC): + a = cutlass.Float32(0.0) + for kk in cutlass.range_constexpr(4): + pos = t - 3 + kk + xv = win[(j + kk) % 4, v] + if pos < s0: + xv = cutlass.Float32(0.0) + if init != 0: + if slot != pad_slot_id: + xv = mState[slot, c0 + v, 3 + pos - s0].to(cutlass.Float32) + a = _fma_rn(xv, w[v, kk], a) + acc[v] = a + y0, y1 = _silu2(acc[0], acc[1]) + y2, y3 = _silu2(acc[2], acc[3]) + p01 = _bf16x2(y0, y1) + p23 = _bf16x2(y2, y3) + if is_qk: + # L2 norm of the rounded head row: fp32 squares, a 32-lane butterfly. + f0, f1 = _unpack2(p01) + f2, f3 = _unpack2(p23) + s_a, s_b = _ffma2(f0, f1, f0, f1, cutlass.Float32(0.0), cutlass.Float32(0.0)) + s_a, s_b = _ffma2(f2, f3, f2, f3, s_a, s_b) + ss = s_a + s_b + for off in cutlass.range_constexpr(5): + ss = ss + cute.arch.shuffle_sync_bfly(ss, offset=1 << off) + inv = _inv_norm(ss, eps) + n0, n1 = _fmul2(f0, f1, inv, inv) + n2, n3 = _fmul2(f2, f3, inv, inv) + p01 = _bf16x2(n0, n1) + p23 = _bf16x2(n2, n3) + if t < num_tokens: + _st_v2(o_base + cutlass.Int64(t) * o_stride, p01, p23) + + +@triton.jit +def _gdn_gates_kernel( + a_ptr, + stride_at, + stride_ah, + b_ptr, + stride_bt, + stride_bh, + A_log_ptr, + dt_bias_ptr, + g_ptr, + beta_ptr, + num_tokens, + softplus_threshold, + NUM_V_HEADS: tl.constexpr, + HP: tl.constexpr, + BLOCK_T: tl.constexpr, +): + """The decay/beta gates over `[T, NUM_V_HEADS]`, one row-coalesced tile a program: + the Triton conv pass's code for them.""" + t = (tl.program_id(0) * BLOCK_T + tl.arange(0, BLOCK_T))[:, None] + h = tl.arange(0, HP)[None, :] + m = (t < num_tokens) & (h < NUM_V_HEADS) + t64 = t.to(tl.int64) + a_val = tl.load(a_ptr + t64 * stride_at + h * stride_ah, mask=m, other=0.0) + b_val = tl.load(b_ptr + t64 * stride_bt + h * stride_bh, mask=m, other=0.0) + A_log = tl.load(A_log_ptr + h, mask=h < NUM_V_HEADS, other=0.0).to(tl.float32) + dt_bias = tl.load(dt_bias_ptr + h, mask=h < NUM_V_HEADS, other=0.0).to(tl.float32) + gate_input = a_val.to(tl.float32) + dt_bias + softplus = tl.where( + gate_input <= softplus_threshold, tl.log(1.0 + tl.exp(gate_input)), gate_input + ) + g_off = t64 * NUM_V_HEADS + h + tl.store(g_ptr + g_off, -tl.exp(A_log) * softplus, mask=m) + tl.store(beta_ptr + g_off, tl.sigmoid(b_val.to(tl.float32)), mask=m) + + +_CONV_PASS: dict = {} + + +def _cute_conv_pass_ok(x, weight, conv_states, num_k_heads, head_k_dim, num_v_heads) -> bool: + """The CuTe conv pass's preconditions: bf16, 128-channel heads, 16-byte aligned rows + of `x` (its `cp.async` copies), int32 element strides.""" + return ( + x.dtype == torch.bfloat16 + and head_k_dim == _VEC * 32 + and x.data_ptr() % 16 == 0 + and (x.stride(0) * 2) % 16 == 0 + and max(*weight.stride(), *conv_states.stride()) < 2**31 + ) + + +def _cute_conv_pass( + x, weight, conv_states, query_start_loc, cache_indices, has_initial_state, q, k, v, + num_k_heads, num_v_heads, head_k_dim, pad_slot_id, l2_norm_eps, +): # fmt: skip + num_tokens = x.shape[0] + num_seqs = query_start_loc.shape[0] - 1 + key = (num_k_heads, num_v_heads, head_k_dim, x.stride(0)) + q2 = q.view(num_tokens, -1) + ptrs = ( + x.data_ptr(), + weight.data_ptr(), + conv_states.data_ptr(), + query_start_loc.data_ptr(), + cache_indices.data_ptr(), + has_initial_state.data_ptr(), + q2, + k.data_ptr(), + v.data_ptr(), + ) + fn = _CONV_PASS.get(key) + if fn is None: + gm = cute.AddressSpace.gmem + types = ( + (cutlass.BFloat16, 16), + (cutlass.BFloat16, 2), + (cutlass.BFloat16, 2), + (cutlass.Int32, 4), + (cutlass.Int32, 4), + (cutlass.Uint8, 1), + None, + (cutlass.BFloat16, 16), + (cutlass.BFloat16, 16), + ) + cptrs = [ + from_dlpack(p, assumed_align=16).mark_layout_dynamic(leading_dim=1) + if spec is None + else make_ptr(spec[0], p, gm, assumed_align=spec[1]) + for spec, p in zip(types, ptrs) + ] + fn = cute.compile( + _ConvPassKernel(num_k_heads, num_v_heads, head_k_dim, x.stride(0)), + *cptrs, + cutlass.Int32(1), + cutlass.Int32(1), + cutlass.Int32(1), + cutlass.Int32(-1), + cutlass.Float32(1e-6), + *(cutlass.Int32(1),) * 5, + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--opt-level 2 --enable-tvm-ffi", + ) + _CONV_PASS[key] = fn + fn( + *ptrs, + num_tokens, + num_seqs, + max(num_seqs - 1, 0).bit_length() + 1, + pad_slot_id, + l2_norm_eps, + *weight.stride(), + *conv_states.stride(), + ) + + +def _launch( + x, + weight, + conv_states, + query_start_loc, + cache_indices, + has_initial_state, + a, + b, + A_log, + dt_bias, + num_k_heads, + head_k_dim, + num_v_heads, + head_v_dim, + pad_slot_id, + l2_norm_eps, + softplus_threshold, + block_t=_BLOCK_T, + num_warps=_NUM_WARPS, + maxnreg=_MAXNREG, +): + assert head_k_dim == head_v_dim and weight.shape[-1] == CONV_WIDTH, ( + "one program per 128-channel head: head_k_dim must equal head_v_dim, and the " + f"conv width must be {CONV_WIDTH}" + ) + num_tokens = x.shape[0] + num_seqs = query_start_loc.shape[0] - 1 + conv_dim = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim + dev, dt = x.device, x.dtype + q = torch.empty((1, num_tokens, num_k_heads, head_k_dim), dtype=dt, device=dev) + k = torch.empty_like(q) + v = torch.empty((1, num_tokens, num_v_heads, head_v_dim), dtype=dt, device=dev) + g = torch.empty((1, num_tokens, num_v_heads), dtype=torch.float32, device=dev) + beta = torch.empty((1, num_tokens, num_v_heads), dtype=torch.float32, device=dev) + if num_seqs == 0: + return q, k, v, g, beta + if num_tokens > 0 and _cute_conv_pass_ok( + x, weight, conv_states, num_k_heads, head_k_dim, num_v_heads + ): + _cute_conv_pass( + x, weight, conv_states, query_start_loc, cache_indices, has_initial_state, q, k, v, + num_k_heads, num_v_heads, head_k_dim, pad_slot_id, l2_norm_eps, + ) # fmt: skip + _gdn_gates_kernel[(triton.cdiv(num_tokens, 64),)]( + a, + a.stride(0), + a.stride(1), + b, + b.stride(0), + b.stride(1), + A_log, + dt_bias, + g, + beta, + num_tokens, + softplus_threshold, + NUM_V_HEADS=num_v_heads, + HP=triton.next_power_of_2(num_v_heads), + BLOCK_T=64, + num_warps=4, + ) + elif num_tokens > 0: + heads = 2 * num_k_heads + num_v_heads + blocks = triton.cdiv(num_tokens, block_t) + # int32 element offsets suffice while the largest one fits. + out_width = max(num_k_heads * head_k_dim, num_v_heads * head_v_dim) + off64 = (num_tokens + 3) * max(x.stride(0), a.stride(0), b.stride(0), out_width) >= 2**31 + _gdn_prefill_conv_post_kernel[(heads, blocks)]( + x, + x.stride(0), + weight, + weight.stride(0), + weight.stride(1), + conv_states, + *conv_states.stride(), + query_start_loc, + cache_indices, + has_initial_state, + a, + a.stride(0), + a.stride(1), + b, + b.stride(0), + b.stride(1), + A_log, + dt_bias, + q, + k, + v, + g, + beta, + num_tokens, + num_seqs, + max(num_seqs - 1, 0).bit_length() + 1, + pad_slot_id, + l2_norm_eps, + softplus_threshold, + IDX=tl.int64 if off64 else tl.int32, + NUM_K_HEADS=num_k_heads, + NUM_V_HEADS=num_v_heads, + HEAD_DIM=head_k_dim, + BLOCK_T=block_t, + num_warps=num_warps, + maxnreg=maxnreg, + ) + _gdn_prefill_conv_publish_kernel[(num_seqs, triton.cdiv(conv_dim, _PUBLISH_BLOCK_C))]( + x, + x.stride(0), + conv_states, + *conv_states.stride(), + query_start_loc, + cache_indices, + has_initial_state, + conv_dim, + pad_slot_id, + BLOCK_C=_PUBLISH_BLOCK_C, + num_warps=4, + ) + return q, k, v, g, beta + + +def _engine_slots(num_seqs: int) -> tuple[int, ...]: + """`num_seqs` distinct slots of the engine pool, the last and slot 0 first.""" + rest = tuple((i * 37 + 4) % (ENGINE_SLOTS - 2) + 1 for i in range(max(num_seqs - 2, 0))) + slots = ((ENGINE_SLOTS - 1, 0) + rest)[:num_seqs] + assert len(set(slots)) == num_seqs + return slots + + +class GdnPrefillConvPost(OpWrapper): + """Causal conv + SiLU of the context rows, then the GDN q/k/v/g/beta. + + `x` is `[T, >= qkv_dim]` token-major with unit channel stride (any row + stride: the target passes the q|k|v column view of its `[T, 16480]` input + projection), packed `[q | k | v]`, `qkv_dim = 2 * num_k_heads * head_k_dim + + num_v_heads * head_v_dim`. For sequence `s` over rows `[b, e)` of + `query_start_loc` and `slot = cache_indices[s]`: + + xpad = cat([conv_states[slot] if has_initial_state[s] else 0, x[b:e]]) + y[t] = bf16(silu(sum_k weight[:, k] * xpad[t + k])) (fp32 taps, in order) + q, k = y / sqrt(sum(y^2) + l2_norm_eps) (per head, fp32, rounded once) + v = y (bit-exact) + g = -exp(A_log) * softplus(a + dt_bias) (fp32) + beta = sigmoid(b) (fp32) + conv_states[slot] = xpad[e - b : e - b + W - 1] (raw input, exact) + + -- the composition `extract_transpose_prefill_slice` -> `causal_conv1d_fwd` + (SiLU, no bias) -> `fused_gdn_post_conv(prefill, None, ...)`, in one pass + over the activation plus a per-sequence window publish. + + Fusion boundary: the conv-state read, conv, SiLU, split, q/k norm, gates + and the state publish are inside. The projection, the generation rows' + conv and the recurrence are the caller's (`use_qk_l2norm_in_kernel=False` + downstream). + + Returns `(q, k, v, g, beta)`: fresh contiguous `[1, T, num_k_heads, + head_k_dim]` x2, `[1, T, num_v_heads, head_v_dim]` in `x.dtype`, and + `[1, T, num_v_heads]` fp32 twice. `conv_states` is advanced in place for + every sequence whose slot is not `pad_slot_id`; nothing else is written. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: lengths; slots (pool rows used, None -> _engine_slots); pool_slots; + # initial; layout "in_proj" (x and a/b column views of one [T, 16480] + # buffer, the target's call) | "contiguous" (x [T, 10240], a/b [T, 48]) | + # "rows" (x/a/b the first rows of a longer projection, as a mixed step's + # context half); pad_slot_id; dt_bias_fill. + CELLS: tuple[Cell, ...] = ( + Cell( + why="target call: 37+1+128 tokens, in_proj views, 8-slot pool", + spec=dict( + lengths=(37, 1, 128), pool_slots=8, slots=(3, 0, 5), initial=(True, False, True) + ), + ), + Cell( + why="engine 461-slot pool: slots 460, 0, 230", + spec=dict( + lengths=(37, 1, 128), + pool_slots=ENGINE_SLOTS, + slots=(ENGINE_SLOTS - 1, 0, ENGINE_SLOTS // 2), + initial=(True, False, True), + ), + ), + Cell( + why="one 8192-token context chunk with carry", + spec=dict(lengths=(8192,), pool_slots=ENGINE_SLOTS, initial=(True,)), + ), + Cell( + why="max_num_tokens 16384 over 3 contexts (the con=128 mixed step)", + spec=dict( + lengths=(5440, 5440, 5504), pool_slots=ENGINE_SLOTS, initial=(True, False, True) + ), + ), + Cell( + why="mixed step's context half: first 8065 rows of a 8192-row projection", + spec=dict( + lengths=(8065,), + pool_slots=ENGINE_SLOTS, + initial=(False,), + layout="rows", + extra_rows=127, + ), + ), + Cell( + why="128 short contexts (the batch ceiling) over the 461-slot pool", + spec=dict( + lengths=tuple(1 + (i * 7) % 64 for i in range(128)), + pool_slots=ENGINE_SLOTS, + initial=tuple(i % 3 != 0 for i in range(128)), + ), + ), + Cell( + why="contexts shorter than the window: carry part shifted, part tokens", + spec=dict( + lengths=(2, 1, 2, 3), + pool_slots=6, + slots=(0, 1, 2, 5), + initial=(True, True, False, True), + ), + ), + Cell( + why="a zero-token context republishes its carry unchanged", + spec=dict( + lengths=(0, 40, 0), pool_slots=4, slots=(2, 3, 1), initial=(True, False, False) + ), + ), + Cell( + why="contiguous x [T, 10240] and a/b [T, 48]", + spec=dict( + lengths=(64, 9), + pool_slots=4, + slots=(1, 2), + initial=(True, True), + layout="contiguous", + ), + ), + Cell( + why="pad_slot_id -1: that context neither reads nor writes the pool", + spec=dict( + lengths=(16, 8, 16), pool_slots=4, slots=(0, -1, 2), initial=(True, True, True) + ), + ), + Cell( + why="every head below softplus_threshold", + spec=dict( + lengths=(64,), + pool_slots=2, + slots=(1,), + initial=(True,), + dt_bias_fill=SOFTPLUS_THRESHOLD - 8.0, + ), + ), + Cell( + why="every head above softplus_threshold: the linear branch", + spec=dict( + lengths=(64,), + pool_slots=2, + slots=(1,), + initial=(True,), + dt_bias_fill=SOFTPLUS_THRESHOLD + 8.0, + ), + ), + ) + + note = """ + Metadata consumed: none. Batch layout (`query_start_loc`, int32 `[N + 1]`, + `query_start_loc[0] == 0`, `query_start_loc[N] == x.shape[0]`), slot map + (`cache_indices`, int32 `[N]`) and carry flags (`has_initial_state`, bool + `[N]`) are explicit. + + Call sites (`PrefillTarget._gdn_prefill_prep`, from `_gdn_block`): the + context rows of every prefill-carrying step -- prefill-only `(mixed_qkv, + a, b, qsl_all, si_all, hi_all)`, mixed and MTP-verify `(mq_p, a_p, b_p, + qsl_p, si_p, hi_p)` -- with the layer's `[10240, 4]` conv weight, its + whole conv pool unsliced, `pad_slot_id` -1, `l2_norm_eps` 1e-6 and + `softplus_threshold` 20. `x`, `a`, `b` are column (and row) views, row + stride 16480, of the one input-projection output; the outputs feed + `flashinfer_chunk_gated_delta_rule` (the Triton chunk on the verify path) + with `use_qk_l2norm_in_kernel=False`. + + Three launches. The conv pass (CuTe DSL, `_ConvPassKernel`): grid + `80 heads x cdiv(T, 64)`, 128 threads; a CTA copies its 67-row tile + (3 history rows) of the head's 128 channels into shared memory with + 16-byte `cp.async`, then a thread slides the width-4 window over 16 + tokens of 4 channels, each element converted once; each thread finds its + first token's sequence by a binary search of `query_start_loc` + (sync-free, capture-safe), and a thread whose 16 tokens and windows lie + inside one sequence takes a straight-line branch, the rest read the pool + token by token. The gates (Triton): one row-coalesced `[64, 48]` tile a + program. Then a per-sequence window publish (Triton), after both, so no + program reads a window another has already advanced. Inputs the CuTe + pass cannot take (not bf16, a row of `x` not 16-byte aligned, other head + sizes) run the Triton conv pass (`_gdn_prefill_conv_post_kernel`, gates + included) instead. + + Speed (one layer, 461-slot pool, in_proj views, us, graph replay over 4 + layers' inputs): 16384 tokens / 3 contexts 136 (the Triton conv pass 198, + the three-call chain it replaced 528). The Triton pass is bound by its + instruction stream (ncu: 183 M warp instructions, 74% issue-busy, DRAM + 41%); the CuTe pass's straight-line window and paired fp32 arithmetic + issue a third fewer, and the shared-memory tile keeps the bytes in flight + out of registers. + + Numerics: the conv result is rounded to bf16 before the norm reads it, as + the unfused path stores it. Conv taps: fp32 FMAs from zero in tap order + (`fma.rn.f32x2` on channel pairs). SiLU: `acc * rcp(1 + 2^(-acc * + log2e))` with `ex2.approx.ftz` and `rcp.approx.ftz` -- the Triton pass's + `ex2.approx` / `div.full` values wherever `1 + e <= 2^126`; below + acc = -87.3 it returns -0 where those return a value under 1e-36. q/k + norm: the 128-channel fp32 sum of squares as each thread's pairs plus a + 32-lane butterfly, `sqrt.approx.ftz`, `div.full`. v, g, beta and the pool + are bitwise the Triton pass's; q/k one bf16 ulp on ~3e-6 of elements + (the sum's order). Against the unfused path (entry test, every cell): + q/k/v within one bf16 ulp on 1e-6..1e-4 of elements, g/beta within the + fp32 band, conv pool bitwise. + + A `pad_slot_id` context reads no carry and publishes nothing; its outputs + are computed from a zero window (the unfused conv leaves those columns + unconvolved instead -- nothing consumes either). The engine gives every + context row a real slot. + + Preconditions (`is_valid`): `x` 2-D, unit channel stride, at least + `qkv_dim` wide; `a`/`b` at least `T` rows (read through both strides); + A_log/dt_bias contiguous with `num_v_heads` entries; `weight` `[qkv_dim, + 4]` in `x`'s dtype (any strides); `conv_states` `[slots, qkv_dim, 3]` in + `x`'s dtype (any strides); int32 `query_start_loc`/`cache_indices`, bool + `has_initial_state`. Loud (`__call__`): `head_k_dim != head_v_dim`, conv + width other than 4. + + Not checked: `cache_indices` within the pool and distinct; + `query_start_loc` non-decreasing and ending at `T`. + + Not certified: fp16/fp32 activations, other head geometries, bias. + + Hard limits + element offsets int32 unless the call needs int64 (chosen on the host + from T and the strides) not probed beyond 16384 tokens. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + conv_states: torch.Tensor, + query_start_loc: torch.Tensor, + cache_indices: torch.Tensor, + has_initial_state: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + num_k_heads: int, + head_k_dim: int, + num_v_heads: int, + head_v_dim: int, + pad_slot_id: int = -1, + l2_norm_eps: float = 1e-6, + softplus_threshold: float = 20.0, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + return _launch( + x, + weight, + conv_states, + query_start_loc, + cache_indices, + has_initial_state, + a, + b, + A_log, + dt_bias, + num_k_heads, + head_k_dim, + num_v_heads, + head_v_dim, + pad_slot_id, + l2_norm_eps, + softplus_threshold, + ) + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + conv_states: torch.Tensor, + query_start_loc: torch.Tensor, + cache_indices: torch.Tensor, + has_initial_state: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + num_k_heads: int, + head_k_dim: int, + num_v_heads: int, + head_v_dim: int, + pad_slot_id: int = -1, + l2_norm_eps: float = 1e-6, + softplus_threshold: float = 20.0, + ) -> tuple[torch.Tensor, ...]: + """`(q, k, v, g, beta, conv_states as the call leaves them)`; mutates + nothing. Conv taps in fp32 in tap order and the SiLU in fp32, rounded + to `x.dtype` once; the q/k norm in fp32 from that rounded value, + rounded once; `g` in fp64 with the stabilised softplus, rounded to + fp32; `beta` the fp32 sigmoid.""" + qkv_dim = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim + width = weight.shape[-1] + n_tok = x.shape[0] + w = weight.float() + y = torch.zeros(n_tok, qkv_dim, dtype=x.dtype, device=x.device) + states = conv_states.clone() + bounds = query_start_loc.tolist() + slots = cache_indices.tolist() + carry = has_initial_state.tolist() + for s in range(len(bounds) - 1): + lo, hi, slot = bounds[s], bounds[s + 1], slots[s] + tokens = x[lo:hi, :qkv_dim].float() + if slot != pad_slot_id and carry[s]: + prefix = conv_states[slot].float().t() + else: + prefix = tokens.new_zeros(width - 1, qkv_dim) + padded = torch.cat([prefix, tokens], dim=0) + length = hi - lo + acc = torch.zeros_like(tokens) + for kk in range(width): + acc = acc + padded[kk : kk + length] * w[:, kk] + y[lo:hi] = torch.nn.functional.silu(acc).to(x.dtype) + if slot != pad_slot_id: + states[slot] = padded[length : length + width - 1].t().to(states.dtype) + kd = num_k_heads * head_k_dim + + def l2norm(part): + part = part.reshape(n_tok, num_k_heads, head_k_dim).float() + return (part / torch.sqrt(part.pow(2).sum(-1, keepdim=True) + l2_norm_eps)).to(x.dtype) + + q, k = l2norm(y[:, :kd]), l2norm(y[:, kd : 2 * kd]) + v = y[:, 2 * kd :].reshape(n_tok, num_v_heads, head_v_dim) + gx = a[:n_tok, :num_v_heads].double() + dt_bias[:num_v_heads].double() + softplus = torch.where(gx <= softplus_threshold, torch.log1p(torch.exp(gx)), gx) + g = (-torch.exp(A_log[:num_v_heads].double()) * softplus).float() + beta = torch.sigmoid(b[:n_tok, :num_v_heads].float()) + return (*(t.unsqueeze(0) for t in (q, k, v, g, beta)), states) + + def is_valid( + self, + x: torch.Tensor, + weight: torch.Tensor, + conv_states: torch.Tensor, + query_start_loc: torch.Tensor, + cache_indices: torch.Tensor, + has_initial_state: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + num_k_heads: int, + head_k_dim: int, + num_v_heads: int, + head_v_dim: int, + pad_slot_id: int = -1, + l2_norm_eps: float = 1e-6, + softplus_threshold: float = 20.0, + ) -> None: + # Every address below is computed from the arguments' own extents and + # strides only where stated; the rest is `ptr + index`, unmasked. + qkv_dim = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim + n_tok = x.shape[0] + assert x.dim() == 2 and x.stride(1) == 1 and x.shape[1] >= qkv_dim, ( + f"x must be [T, >= {qkv_dim}] with unit channel stride; got {tuple(x.shape)} " + f"strides {x.stride()} -- the channel axis is read as `ptr + c`" + ) + assert weight.dtype == x.dtype and weight.shape[0] == qkv_dim, ( + f"weight must be [{qkv_dim}, 4] in x's dtype; got {tuple(weight.shape)} {weight.dtype}" + ) + assert conv_states.dtype == x.dtype and tuple(conv_states.shape[1:]) == (qkv_dim, 3), ( + f"conv_states must be [slots, {qkv_dim}, 3] in x's dtype; got " + f"{tuple(conv_states.shape)} {conv_states.dtype}" + ) + assert a.shape[0] >= n_tok and b.shape[0] >= n_tok, ( + f"a and b must have at least T={n_tok} rows; the token axis is unmasked" + ) + assert A_log.stride(0) == 1 and dt_bias.stride(0) == 1, ( + "A_log and dt_bias must be contiguous 1-D; they are read as `ptr + head`" + ) + assert A_log.numel() >= num_v_heads and dt_bias.numel() >= num_v_heads, ( + f"A_log and dt_bias must hold num_v_heads={num_v_heads} entries; the load is unmasked" + ) + assert query_start_loc.dtype == torch.int32 and cache_indices.dtype == torch.int32, ( + "query_start_loc and cache_indices must be int32" + ) + assert has_initial_state.dtype == torch.bool, "has_initial_state must be bool" + n_seq = query_start_loc.shape[0] - 1 + assert cache_indices.shape[0] >= n_seq and has_initial_state.shape[0] >= n_seq, ( + f"cache_indices and has_initial_state must cover the {n_seq} sequences" + ) + + +gdn_prefill_conv_post = GdnPrefillConvPost() diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_rms_norm_gated_token_major.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_rms_norm_gated_token_major.py index 381ef675e1fe..f3a733751962 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_rms_norm_gated_token_major.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_rms_norm_gated_token_major.py @@ -31,6 +31,8 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: #: The GDN in_proj row the old column-slice measurement used: [q k v | z] = 10240 | 6144. _PROJ_WIDTH, _Z_OFFSET = 16384, 10240 +#: The target's concatenated GDN input projection: [q k v | z | a | b] = 10240 | 6144 | 48 | 48. +_IN_PROJ_WIDTH = 16480 def _seed(text: str) -> int: @@ -55,6 +57,8 @@ def randn(*shape): 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 == "in_proj": + z = randn(t, _IN_PROJ_WIDTH)[:, _Z_OFFSET : _Z_OFFSET + heads * n].view(t, heads, n) elif layout == "head_strided": z = randn(t, heads, 2 * n)[:, :, :n] else: @@ -76,6 +80,8 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res 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("z") == "in_proj" and cell.spec["tokens"] > 1: + assert z.stride() == (_IN_PROJ_WIDTH, cell.spec["width"], 1) 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()] diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_triton_update.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_triton_update.py index c9ef151454bf..110c24a273ed 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_triton_update.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_triton_update.py @@ -61,6 +61,12 @@ def randn(*shape): x = randn(batch, dim) elif x_form == "contiguous": x = randn(batch, dim, seqlen) + elif x_form == "in_proj": + # The call site under the one input GEMM: the q|k|v columns of an + # [n * p, C + 6240] projection, reshaped [n, p, C] and seen as [n, C, p]. + wide = dim + 6144 + 2 * 48 + x = randn(batch * seqlen, wide)[:, :dim].reshape(batch, seqlen, dim).transpose(1, 2) + assert x.stride() == (seqlen * wide, 1, wide) 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) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_update.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_update.py index 4b61e62ec676..9f468a8640b9 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_update.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_update.py @@ -41,6 +41,8 @@ def _compare( PAD_SLOT_ID = -1 #: Rows of prefill ahead of the decode rows in the mixed-batch cell. _PREFILL_ROWS = 37 +#: Columns past q|k|v in the target's one GDN input-projection output: z | a | b. +_IN_PROJ_TAIL = 6144 + 2 * 48 def _build(spec, seed): @@ -59,6 +61,10 @@ def randn(*shape): elif x_form == "2d_row_slice": extras["buf"] = randn(_PREFILL_ROWS + batch, dim) x = extras["buf"][_PREFILL_ROWS:] + elif x_form == "2d_in_proj": + extras["buf"] = randn(_PREFILL_ROWS + batch, dim + _IN_PROJ_TAIL) + x = extras["buf"][_PREFILL_ROWS:, :dim] + assert x.stride() == (dim + _IN_PROJ_TAIL, 1) elif x_form == "token_major": x = randn(batch, seqlen, dim).transpose(1, 2) assert x.stride() == (seqlen * dim, 1, dim) @@ -91,7 +97,7 @@ def _run(args, extras): got = list(args) if "buf" in extras: got_buf = extras["buf"].clone() - got[0] = got_buf[_PREFILL_ROWS:] + got[0] = got_buf[_PREFILL_ROWS:, : args[0].shape[1]] else: got[0] = args[0].clone() if "wide" in extras: @@ -103,6 +109,8 @@ def _run(args, extras): 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]) + dim = args[0].shape[1] + assert torch.equal(got_buf[:, dim:], extras["buf"][:, dim:]), "columns past x written" if "wide" in extras: assert torch.equal(got_wide[:, :, 1::2], extras["wide"][:, :, 1::2]) return returned, got[1] @@ -179,3 +187,27 @@ def test_op_rejects_every_metadata_violation_loudly(case: str) -> None: with pytest.raises(RuntimeError, match=fragment.replace("(", r"\(").replace(")", r"\)")): with validating(op): op(*args) + + +@pytest.mark.parametrize("rows", [64, 121, 128]) +def test_decode_rows_match_the_trtllm_kernel_bitwise(rows: int) -> None: + """From 64 rows the decode shape runs the CuTe kernel: x and the whole pool + bitwise what `torch.ops.trtllm.causal_conv1d_update` leaves, a pad row untouched. + x is the q|k|v column view of an in_proj output; the pool a per-layer view of a + two-layer buffer, as the engine hands both over.""" + g = torch.Generator(device="cuda").manual_seed(rows) + dim = 10240 + proj = torch.randn(rows, dim + _IN_PROJ_TAIL, generator=g, device="cuda").to(torch.bfloat16) + pool = torch.randn(ENGINE_SLOTS, 2, dim, 3, generator=g, device="cuda").to(torch.bfloat16)[:, 1] + weight = (torch.randn(dim, 4, generator=g, device="cuda") * 0.5).to(torch.bfloat16) + slots = (torch.randperm(ENGINE_SLOTS, generator=g, device="cuda")[:rows]).to(torch.int32) + slots[rows // 3] = PAD_SLOT_ID + x_ref, pool_ref = proj.clone(), pool.clone() + torch.ops.trtllm.causal_conv1d_update( + x_ref[:, :dim].unsqueeze(-1), pool_ref, weight, None, True, None, slots, PAD_SLOT_ID + ) + x_got, pool_got = proj.clone(), pool.clone() + op(x_got[:, :dim], pool_got, weight, None, True, None, slots, PAD_SLOT_ID) + assert torch.equal(x_got, x_ref), "x differs from the TRT-LLM kernel" + assert torch.equal(pool_got, pool_ref), "the pool differs from the TRT-LLM kernel" + assert torch.equal(x_got[rows // 3], proj[rows // 3]), "the pad row was written" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_post_conv.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_post_conv.py index 5cdbfe632874..325bb8925910 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_post_conv.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_post_conv.py @@ -78,6 +78,8 @@ def _compare(out: tuple, ref: tuple) -> None: #: 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 +#: The target's one GDN input-projection output: [q|k|v | z | a | b]. +IN_PROJ_WIDTH = QKV_DIM + Z_WIDTH + 2 * NUM_V_HEADS # 16480 def _build(spec, seed): @@ -96,6 +98,18 @@ def randn(*shape): 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 == "in_proj": + # `_gdn_block`: one [T, 16480] GEMM output; q|k|v, a and b are its column + # views, the decode half the row split of q|k|v, the prefill half the + # transposing extraction of the rest. + proj = randn(n, IN_PROJ_WIDTH) + mixed = proj[:, :QKV_DIM] + prefill = mixed[:n_p].t().contiguous() + decode = mixed[n_p:] if n_d else None + a = proj[:, QKV_DIM + Z_WIDTH : QKV_DIM + Z_WIDTH + NUM_V_HEADS] + b = proj[:, QKV_DIM + Z_WIDTH + NUM_V_HEADS :] + assert a.stride(0) == b.stride(0) == IN_PROJ_WIDTH + assert decode is None or decode.stride() == (IN_PROJ_WIDTH, 1) 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. diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_sigmoid_gating_delta_rule_update.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_sigmoid_gating_delta_rule_update.py index 313b228ff644..4f7a519ce50b 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_sigmoid_gating_delta_rule_update.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_sigmoid_gating_delta_rule_update.py @@ -61,6 +61,9 @@ def _compare( DEV = "cuda" NUM_K_HEADS, NUM_V_HEADS, HEAD_K_DIM, HEAD_V_DIM = 16, 48, 128, 128 MIXED_QKV_WIDTH = 2 * NUM_K_HEADS * HEAD_K_DIM + NUM_V_HEADS * HEAD_V_DIM # 10240 +#: The target's one GDN input-projection output: [q|k|v | z | a | b], 16480 wide. +IN_PROJ_A = MIXED_QKV_WIDTH + NUM_V_HEADS * HEAD_V_DIM # 16384 +IN_PROJ_WIDTH = IN_PROJ_A + 2 * NUM_V_HEADS SOFTPLUS_BETA, SOFTPLUS_THRESHOLD = 1.0, 20.0 _LAYERS = 3 _TAIL = 5 @@ -84,7 +87,11 @@ def _build(spec, seed): seq_lens = spec["seq_lens"] total = sum(seq_lens) layout = spec.get("layout", "varlen") - mixed = torch.randn(total, MIXED_QKV_WIDTH, generator=g, device=DEV, dtype=torch.bfloat16) + if spec.get("ab") == "in_proj": + proj = torch.randn(total, IN_PROJ_WIDTH, generator=g, device=DEV, dtype=torch.bfloat16) + mixed = proj[:, :MIXED_QKV_WIDTH] + else: + mixed = torch.randn(total, MIXED_QKV_WIDTH, generator=g, device=DEV, dtype=torch.bfloat16) kd = NUM_K_HEADS * HEAD_K_DIM q = mixed[:, :kd].view(1, total, NUM_K_HEADS, HEAD_K_DIM) k = mixed[:, kd : 2 * kd].view(1, total, NUM_K_HEADS, HEAD_K_DIM) @@ -93,7 +100,11 @@ def _build(spec, seed): for t in (q, k): tf = t.float() t.copy_(tf / (tf.pow(2).sum(-1, keepdim=True).sqrt() + 1e-6)) - if spec.get("ab", "packed") == "separate": + if spec.get("ab") == "in_proj": + a = proj[:, IN_PROJ_A : IN_PROJ_A + NUM_V_HEADS] + b = proj[:, IN_PROJ_A + NUM_V_HEADS :] + assert mixed.stride(0) == a.stride(0) == b.stride(0) == IN_PROJ_WIDTH + elif 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: diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_gated_norm_pdl.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_gated_norm_pdl.py new file mode 100644 index 000000000000..c3718a602f32 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_gated_norm_pdl.py @@ -0,0 +1,179 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gdn_decode_gated_norm_pdl entry. + +Cells against the composition's fp64 reference (the recurrence entry's, then the +gated norm entry's) in the recurrence entry's band: torch's default for the bf16 +output and the fp32 pool, the pool compared whole so slots the call did not name are +gated too. Beyond the cells: agreement with the unfused composition it replaces +(FlashInfer's recurrence, then the gated norm entry), PDL on and off bitwise, a +stale-read check under graph replay (a producer rewrites the projection inside the +graph right before the op), and the guards. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.gated_delta_rule_decode_pretranspose import ( # noqa: E501 # noqa: E501 + _scattered, + gated_delta_rule_decode_pretranspose, +) +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, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_decode_gated_norm_pdl import ( # noqa: E501 + gdn_decode_gated_norm_pdl as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_decode_gated_norm_pdl import ( # noqa: E501 + launch, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "gdn_decode_gated_norm_pdl requires a CUDA device" + +DEV = "cuda" +H, HV, K, V = 16, 48, 128, 128 +SLOTS, EPS = 461, 1e-6 +WIDTH = 2 * H * K + 2 * HV * V + 2 * HV + + +def _build(spec, seed): + g = torch.Generator(device=DEV).manual_seed(seed) + rows = spec["rows"] + proj = torch.randn(rows, WIDTH, generator=g, device=DEV).to(torch.bfloat16) + if spec.get("pool") == "coalesced": + # The KV-cache-v2 manager's view: every GDN layer's state of a slot in one page, + # this layer's state one sub-page of it, so the slot stride is 36 layers' states + # and slot 100 sits ~2.8G elements past the base (past int32). + layers, slots = 36, 102 + page = torch.empty(slots * layers * HV * V * K, device=DEV) + page[: HV * V * K * 4].normal_(0.0, 0.05, generator=g) + pool = page.as_strided((slots, HV, V, K), (layers * HV * V * K, V * K, K, 1)) + for s in (100, 101): + pool[s].copy_(torch.randn(HV, V, K, generator=g, device=DEV) * 0.05) + elif spec.get("pool") == "layer_view": + cache = torch.randn(3, SLOTS, HV, V, K, generator=g, device=DEV) * 0.05 + pool = cache[1] + else: + pool = torch.randn(SLOTS, HV, V, K, generator=g, device=DEV) * 0.05 + if spec.get("pool") == "coalesced": + slots = torch.tensor([100, 101][:rows], dtype=torch.int32, device=DEV) + else: + slots = torch.tensor(_scattered(rows, SLOTS), dtype=torch.int32, device=DEV) + if spec.get("pad"): + slots[1] = -1 + A_log = torch.empty(HV, device=DEV).uniform_(-5.5, 1.0, generator=g) + dt_bias = torch.randn(HV, generator=g, device=DEV) * 4.0 + w = (1.0 + 0.2 * torch.randn(V, generator=g, device=DEV)).to(torch.bfloat16) + return [proj, pool, slots, A_log, dt_bias, w, EPS, H, K] + + +def _compare(out, pool, ref_out, ref_pool) -> None: + assert out.dtype == torch.bfloat16 and out.shape == ref_out.shape and out.is_contiguous() + assert torch.isfinite(out).all() + torch.testing.assert_close(out, ref_out) + torch.testing.assert_close(pool, ref_pool) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) + ref_out, ref_pool = op.reference(*args) + out = op(*args) + _compare(out, args[1], ref_out, ref_pool) + if cell.spec.get("pad"): + assert (out[1] == 0).all(), "a pad row's output is zero" + + +def _unfused(proj, pool, slots, A_log, dt_bias, w, eps, h, k): + rows = proj.shape[0] + q, kk, v, z, a, b = torch.split(proj, [h * k, h * k, HV * V, HV * V, HV, HV], 1) + o = torch.empty(rows, 1, HV, V, dtype=torch.bfloat16, device=DEV) + gated_delta_rule_decode_pretranspose( + q=q.reshape(rows, 1, h, k), + k=kk.reshape(rows, 1, h, k), + v=v.reshape(rows, 1, HV, V), + state=None, + A_log=A_log, + a=a.reshape(rows, 1, HV), + dt_bias=dt_bias, + b=b.reshape(rows, 1, HV), + scale=None, + output=o, + use_qk_l2norm=True, + initial_state=pool, + initial_state_indices=slots, + output_state_indices=None, + ) + return rms_norm_gated_token_major( + o.reshape(-1, V), z.reshape(rows, HV, V), w, eps, None, "sigmoid" + ).reshape(rows, HV * V) + + +@pytest.mark.parametrize("rows", [1, 4]) +def test_matches_the_unfused_composition(rows) -> None: + args = _build(dict(rows=rows), seed=900 + rows) + pool_unfused = args[1].clone() + want = _unfused(args[0], pool_unfused, *args[2:]) + got = op(*args) + torch.testing.assert_close(got, want) + torch.testing.assert_close(args[1], pool_unfused) + print( + f"\n rows {rows}: output bitwise the composition's {(got == want).float().mean().item():.4f}" + ) + + +def test_pdl_on_and_off_agree_bitwise() -> None: + args = _build(dict(rows=2), seed=31) + pool0 = args[1].clone() + a = launch(*args, pdl=True) + pool_a = args[1].clone() + args[1].copy_(pool0) + b = launch(*args, pdl=False) + assert torch.equal(a, b) and torch.equal(pool_a, args[1]) + + +def test_reads_after_the_wait_under_graph_replay() -> None: + """A producer rewrites the projection inside the graph right before the op; with PDL + a read of it before `griddepcontrol.wait` would see the previous replay's row. The + pool advances one step per replay, as in decode.""" + args = _build(dict(rows=1), seed=41) + src = args[0].clone() + + def step(): + args[0].copy_(src * 1.0) + return op(*args) + + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + pool0 = args[1].clone() + with torch.cuda.stream(warmup): + step() + torch.cuda.current_stream().wait_stream(warmup) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_out = step() + for trial in range(3): + src.copy_(_build(dict(rows=1), seed=100 + trial)[0]) + before = args[1].clone() + graph.replay() + torch.cuda.synchronize() + after = args[1].clone() + args[1].copy_(before) + want = op(src.clone(), *args[1:]) + assert torch.equal(graph_out, want), f"replay {trial} output" + assert torch.equal(after, args[1]), f"replay {trial} pool" + del pool0 + + +def test_guards() -> None: + args = _build(dict(rows=1), seed=5) + with validating(op), pytest.raises(AssertionError): + op(args[0].expand(5, -1).contiguous(), args[1], args[2].expand(5).contiguous(), *args[3:]) + with validating(op), pytest.raises(AssertionError): + op(args[0], args[1].to(torch.bfloat16), *args[2:]) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py new file mode 100644 index 000000000000..2aba437bc7f9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py @@ -0,0 +1,230 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gdn_decode_rows_gated_norm entry. + +Cells against the composition's fp64 reference (the recurrence entry's, then the +gated norm entry's) in the recurrence entry's band: torch's default for the bf16 +output and the fp32 pool, the pool compared whole so slots the call did not name are +gated too. Beyond the cells: agreement with the unfused composition it replaces +(FlashInfer's recurrence, then the gated norm entry) at decode batch sizes, PDL on +and off bitwise, a stale-read check under graph replay (a producer rewrites the +projection inside the graph right before the op), and the guards. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.gated_delta_rule_decode_pretranspose import ( # noqa: E501 + _scattered, + gated_delta_rule_decode_pretranspose, +) +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, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_decode_rows_gated_norm import ( # noqa: E501 + gdn_decode_rows as rows_op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_decode_rows_gated_norm import ( # noqa: E501 + gdn_decode_rows_gated_norm as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_decode_rows_gated_norm import ( # noqa: E501 + launch, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "gdn_decode_rows_gated_norm requires a CUDA device" + +DEV = "cuda" +H, HV, K, V = 16, 48, 128, 128 +SLOTS, EPS = 461, 1e-6 +WIDTH = 2 * H * K + 2 * HV * V + 2 * HV + + +def _build(spec, seed): + g = torch.Generator(device=DEV).manual_seed(seed) + rows = spec["rows"] + proj = torch.randn(rows, WIDTH, generator=g, device=DEV).to(torch.bfloat16) + if spec.get("pool") == "coalesced": + # The KV-cache-v2 manager's view: every GDN layer's state of a slot in one page, + # this layer's state one sub-page of it, so the slot stride is 36 layers' states + # and slot 100 sits ~2.8G elements past the base (past int32). + layers, slots = 36, 102 + page = torch.empty(slots * layers * HV * V * K, device=DEV) + page[: HV * V * K * 4].normal_(0.0, 0.05, generator=g) + pool = page.as_strided((slots, HV, V, K), (layers * HV * V * K, V * K, K, 1)) + for s in (100, 101): + pool[s].copy_(torch.randn(HV, V, K, generator=g, device=DEV) * 0.05) + elif spec.get("pool") == "layer_view": + cache = torch.randn(3, SLOTS, HV, V, K, generator=g, device=DEV) * 0.05 + pool = cache[1] + else: + pool = torch.randn(SLOTS, HV, V, K, generator=g, device=DEV) * 0.05 + if spec.get("pool") == "coalesced": + slots = torch.tensor([100, 101][:rows], dtype=torch.int32, device=DEV) + else: + slots = torch.tensor(_scattered(rows, SLOTS), dtype=torch.int32, device=DEV) + if spec.get("pad"): + slots[1] = -1 + A_log = torch.empty(HV, device=DEV).uniform_(-5.5, 1.0, generator=g) + dt_bias = torch.randn(HV, generator=g, device=DEV) * 4.0 + w = (1.0 + 0.2 * torch.randn(V, generator=g, device=DEV)).to(torch.bfloat16) + return [proj, pool, slots, A_log, dt_bias, w, EPS, H, K] + + +def _compare(out, pool, ref_out, ref_pool) -> None: + assert out.dtype == torch.bfloat16 and out.shape == ref_out.shape and out.is_contiguous() + assert torch.isfinite(out).all() + torch.testing.assert_close(out, ref_out) + torch.testing.assert_close(pool, ref_pool) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) + ref_out, ref_pool = op.reference(*args) + out = op(*args) + _compare(out, args[1], ref_out, ref_pool) + if cell.spec.get("pad"): + assert (out[1] == 0).all(), "a pad row's output is zero" + + +def _unfused(proj, pool, slots, A_log, dt_bias, w, eps, h, k): + rows = proj.shape[0] + q, kk, v, z, a, b = torch.split(proj, [h * k, h * k, HV * V, HV * V, HV, HV], 1) + o = torch.empty(rows, 1, HV, V, dtype=torch.bfloat16, device=DEV) + gated_delta_rule_decode_pretranspose( + q=q.reshape(rows, 1, h, k), + k=kk.reshape(rows, 1, h, k), + v=v.reshape(rows, 1, HV, V), + state=None, + A_log=A_log, + a=a.reshape(rows, 1, HV), + dt_bias=dt_bias, + b=b.reshape(rows, 1, HV), + scale=None, + output=o, + use_qk_l2norm=True, + initial_state=pool, + initial_state_indices=slots, + output_state_indices=None, + ) + return rms_norm_gated_token_major( + o.reshape(-1, V), z.reshape(rows, HV, V), w, eps, None, "sigmoid" + ).reshape(rows, HV * V) + + +@pytest.mark.parametrize("rows", [3, 64, 128]) +def test_matches_the_unfused_composition(rows) -> None: + args = _build(dict(rows=rows), seed=900 + rows) + pool_unfused = args[1].clone() + want = _unfused(args[0], pool_unfused, *args[2:]) + got = op(*args) + torch.testing.assert_close(got, want) + torch.testing.assert_close(args[1], pool_unfused) + print( + f"\n rows {rows}: output bitwise the composition's {(got == want).float().mean().item():.4f}" + ) + + +def test_pdl_on_and_off_agree_bitwise() -> None: + args = _build(dict(rows=8), seed=31) + pool0 = args[1].clone() + a = launch(*args, pdl=True) + pool_a = args[1].clone() + args[1].copy_(pool0) + b = launch(*args, pdl=False) + assert torch.equal(a, b) and torch.equal(pool_a, args[1]) + + +def test_reads_after_the_wait_under_graph_replay() -> None: + """A producer rewrites the projection inside the graph right before the op; with PDL + a read of it before `griddepcontrol.wait` would see the previous replay's row. The + pool advances one step per replay, as in decode.""" + args = _build(dict(rows=16), seed=41) + src = args[0].clone() + + def step(): + args[0].copy_(src * 1.0) + return op(*args) + + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + step() + torch.cuda.current_stream().wait_stream(warmup) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_out = step() + for trial in range(3): + src.copy_(_build(dict(rows=16), seed=100 + trial)[0]) + before = args[1].clone() + graph.replay() + torch.cuda.synchronize() + after = args[1].clone() + args[1].copy_(before) + want = op(src.clone(), *args[1:]) + assert torch.equal(graph_out, want), f"replay {trial} output" + assert torch.equal(after, args[1]), f"replay {trial} pool" + + +def test_guards() -> None: + args = _build(dict(rows=4), seed=5) + with validating(op), pytest.raises(AssertionError): + op(args[0], args[1].to(torch.bfloat16), *args[2:]) + with validating(op), pytest.raises(AssertionError): + op(args[0], args[1], args[2].long(), *args[3:]) + with validating(op), pytest.raises(AssertionError): + op(args[0][:, :-2], *args[1:]) + + +# ---- gdn_decode_rows: the recurrence alone, into a row range of a step's output ---------- + + +def _rows_args(spec, seed): + """`(args, buf)`: the entry's arguments, `out` rows 5.. of a `[1, rows + 9, HV, V]` buffer + (the step's output, the call's rows a range of it), and that buffer.""" + proj, pool, slots, A_log, dt_bias, _, _, h, k = _build(spec, seed) + rows = proj.shape[0] + g = torch.Generator(device=DEV).manual_seed(seed + 1) + buf = torch.randn(1, rows + 9, HV, V, generator=g, device=DEV).to(torch.bfloat16) + out = buf[0, 5 : 5 + rows].view(rows, HV * V) + return [proj, pool, slots, A_log, dt_bias, h, k, out], buf + + +@pytest.mark.parametrize("cell", rows_op.CELLS, ids=[c.why[:44] for c in rows_op.CELLS]) +def test_rows_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + certifying(rows_op) + args, buf = _rows_args(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) + outside = torch.cat([buf[0, :5], buf[0, 5 + args[0].shape[0] :]]).clone() + ref_o, ref_pool = rows_op.reference(*args) + got = rows_op(*args) + assert got.data_ptr() == args[-1].data_ptr(), "written into out" + torch.testing.assert_close(got, ref_o) + torch.testing.assert_close(args[1], ref_pool) + assert torch.equal(torch.cat([buf[0, :5], buf[0, 5 + args[0].shape[0] :]]), outside), ( + "rows outside the range untouched" + ) + if cell.spec.get("pad"): + assert (got[1] == 0).all(), "a pad row's output is zero" + + +@pytest.mark.parametrize("rows", [3, 121]) +def test_rows_then_gated_norm_matches_the_fused_entry(rows) -> None: + """The recurrence alone, then the gated norm entry: the fused entry's output within its + band, and the same pool bitwise (one kernel's recurrence).""" + args = _build(dict(rows=rows), seed=700 + rows) + pool_fused = args[1].clone() + fused = op(args[0], pool_fused, *args[2:]) + o = torch.empty(rows, HV * V, dtype=torch.bfloat16, device=DEV) + rows_op(args[0], args[1], args[2], args[3], args[4], args[7], args[8], o) + assert torch.equal(args[1], pool_fused), "pool" + z = torch.split(args[0], [2 * H * K, HV * V, HV * V, 2 * HV], 1)[2] + normed = rms_norm_gated_token_major( + o.reshape(-1, V), z.reshape(rows, HV, V), args[5], args[6], None, "sigmoid" + ).reshape(rows, HV * V) + torch.testing.assert_close(normed, fused) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_in_proj_conv1d_pdl.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_in_proj_conv1d_pdl.py new file mode 100644 index 000000000000..a7c5c62612f4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_in_proj_conv1d_pdl.py @@ -0,0 +1,153 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gdn_in_proj_conv1d_pdl entry. + +Cells against the entry's fp32 reference: the output within the band the +decode projections use (4 bf16 ulp element-wise against the row scale, 2 ulp +relative RMS -- the GEMV's summation order is the only difference), the pool +the same, a pad row's slot bitwise untouched. Beyond the cells: agreement with +the unfused composition it replaces (`low_m_bf16_mm` + `causal_conv1d_update`); +two decode steps under CUDA-graph replay, the second reading the window the +first one wrote; PDL on and off bitwise; and the guards. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.low_m_bf16_mm import low_m_bf16_mm +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, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_in_proj_conv1d_pdl import ( # noqa: E501 + MAX_ROWS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_in_proj_conv1d_pdl import ( # noqa: E501 + gdn_in_proj_conv1d_pdl as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import ( # noqa: E402,F401 — path above; fixture by name + ULP, + assert_within_ulp, + certifying, + validating, +) + +assert torch.cuda.is_available(), "gdn_in_proj_conv1d_pdl requires a CUDA device" + +N, K, CONV_DIM, PAD = 16480, 2560, 10240, -1 + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0, ulp=ULP[torch.bfloat16]) + + +def _build(rows, pool, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + x = torch.randn(rows, K, generator=g, device="cuda").to(torch.bfloat16) + weight = (torch.randn(N, K, generator=g, device="cuda") * K**-0.5).to(torch.bfloat16) + state = torch.randn(pool, CONV_DIM, 3, generator=g, device="cuda").to(torch.bfloat16) + taps = (torch.randn(CONV_DIM, 4, generator=g, device="cuda") * 0.3).to(torch.bfloat16) + return x, weight.t(), state, taps + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 + certifying(op) + s = cell.spec + x, mat_b, state, taps = _build(s["rows"], s["pool"], seed=zlib.crc32(cell.why.encode()) % 2**31) + slots = torch.tensor(s["slots"], dtype=torch.int32, device="cuda") + pool = state.clone() + out = op(x, mat_b, pool, taps, slots, CONV_DIM, PAD, s["pdl"]) + ref_out, ref_state = op.reference(x, mat_b, state, taps, slots, CONV_DIM, PAD, s["pdl"]) + assert out.shape == (s["rows"], N) and out.dtype == torch.bfloat16 and out.is_contiguous() + _compare(out, ref_out) + _compare(pool, ref_state) + touched = {v for v in s["slots"] if v != PAD} + for slot in range(s["pool"]): + if slot not in touched: + assert torch.equal(pool[slot], state[slot]), f"slot {slot} was written" + + +@pytest.mark.parametrize("rows", [1, 2]) +def test_matches_the_unfused_composition(rows) -> None: + x, mat_b, state, taps = _build(rows, 461, seed=5 + rows) + slots = torch.tensor([37, 211][:rows], dtype=torch.int32, device="cuda") + fused_pool, unfused_pool = state.clone(), state.clone() + fused = op(x, mat_b, fused_pool, taps, slots, CONV_DIM, PAD) + unfused = low_m_bf16_mm(x, mat_b, False) + causal_conv1d_update( + x=unfused[:, :CONV_DIM], + conv_state=unfused_pool, + weight=taps, + bias=None, + silu_activation=True, + cache_seqlens=None, + conv_state_indices=slots, + pad_slot_id=PAD, + ) + _compare(fused, unfused) + _compare(fused_pool, unfused_pool) + + +@pytest.mark.parametrize("pdl", [False, True]) +def test_two_steps_under_graph_replay(pdl) -> None: + """Decode runs this under CUDA graphs, layer after layer: capture one step, + replay it twice with new activations. The second replay's conv reads the + window the first replay advanced, so a stale or misplaced window write fails + the second comparison.""" + x, mat_b, state, taps = _build(1, 64, seed=3) + slots = torch.tensor([9], dtype=torch.int32, device="cuda") + pool = state.clone() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(2): + op(x, mat_b, pool.clone(), taps, slots, CONV_DIM, PAD, pdl) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = op(x, mat_b, pool, taps, slots, CONV_DIM, PAD, pdl) + want_state = state.clone() + for trial in range(2): + g = torch.Generator(device="cuda").manual_seed(50 + trial) + x.copy_(torch.randn(x.shape, generator=g, device="cuda").to(x.dtype)) + graph.replay() + torch.cuda.synchronize() + ref_out, want_state = op.reference(x, mat_b, want_state, taps, slots, CONV_DIM, PAD) + _compare(out, ref_out) + _compare(pool, want_state) + + +def test_pdl_on_and_off_agree_bitwise() -> None: + x, mat_b, state, taps = _build(2, 461, seed=13) + slots = torch.tensor([3, 400], dtype=torch.int32, device="cuda") + on_pool, off_pool = state.clone(), state.clone() + on = op(x, mat_b, on_pool, taps, slots, CONV_DIM, PAD, True) + off = op(x, mat_b, off_pool, taps, slots, CONV_DIM, PAD, False) + assert torch.equal(on, off) and torch.equal(on_pool, off_pool) + + +def test_band_rejects_a_conv_on_the_wrong_columns() -> None: + """The near miss a column offset would produce: the conv applied one column off.""" + x, mat_b, state, taps = _build(1, 461, seed=17) + slots = torch.tensor([42], dtype=torch.int32, device="cuda") + out = op(x, mat_b, state.clone(), taps, slots, CONV_DIM, PAD) + ref_out, _ = op.reference(x, mat_b, state, taps, slots, CONV_DIM, PAD) + _compare(out, ref_out) + with pytest.raises(AssertionError): + _compare(out, torch.roll(ref_out, 1, dims=1)) + + +def test_guards() -> None: + x, mat_b, state, taps = _build(MAX_ROWS + 1, 16, seed=21) + slots = torch.zeros(MAX_ROWS + 1, dtype=torch.int32, device="cuda") + with validating(op), pytest.raises(AssertionError): + op(x, mat_b, state, taps, slots, CONV_DIM, PAD) + x1 = x[:1].contiguous() + with validating(op), pytest.raises(AssertionError): + op(x1, mat_b.contiguous(), state, taps, slots[:1], CONV_DIM, PAD) + with validating(op), pytest.raises(AssertionError): + op(x1, mat_b, state, taps, slots[:1].long(), CONV_DIM, PAD) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_prefill_conv_post.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_prefill_conv_post.py new file mode 100644 index 000000000000..dafeb0921e6c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_prefill_conv_post.py @@ -0,0 +1,218 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gdn_prefill_conv_post entry. + +Every cell is checked twice: against the entry's fp32/fp64 reference, and +against the three-call composition it replaces in the target +(`extract_transpose_prefill_slice` -> `causal_conv1d_fwd` -> +`fused_gdn_post_conv`), with the conv pool compared bitwise both times. +Geometry is a Qwen3.8-Flash-Next linear-attention layer: 16 x 128 key heads, +48 x 128 value heads, qkv_dim 10240, conv width 4, bf16 activations. +""" + +import math +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.causal_conv1d_fwd import ( # noqa: E501 + causal_conv1d_fwd, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.extract_transpose_prefill_slice import ( # noqa: E501 + extract_transpose_prefill_slice, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_gdn_post_conv import ( # noqa: E501 + fused_gdn_post_conv, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_prefill_conv_post import ( # noqa: E501 + A_LOG_RANGE, + CONV_WIDTH, + DT_BIAS_RANGE, + HEAD_K_DIM, + HEAD_V_DIM, + IN_PROJ_WIDTH, + NUM_K_HEADS, + NUM_V_HEADS, + QKV_DIM, + SOFTPLUS_THRESHOLD, + _engine_slots, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_prefill_conv_post import ( # noqa: E501 + gdn_prefill_conv_post as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "gdn_prefill_conv_post requires a CUDA device" + +DEV = "cuda" +DTYPE = torch.bfloat16 +HEADS = (NUM_K_HEADS, HEAD_K_DIM, NUM_V_HEADS, HEAD_V_DIM) +PAD = -1 +L2_NORM_EPS = 1e-6 +#: `g`'s band, derived as in the fused_gdn_post_conv entry test: two SASS +#: approximate transcendentals (~2^-22 relative each), the softplus bounded by +#: the threshold, every term scaled by exp(A_log). +_APPROX_ULP = 2.0**-22 +G_ATOL = 4 * _APPROX_ULP * SOFTPLUS_THRESHOLD * math.exp(A_LOG_RANGE[1]) +A_COL = QKV_DIM + NUM_V_HEADS * HEAD_V_DIM + + +def _build(spec, seed): + """`__call__`'s positional arguments for one cell.""" + g = torch.Generator(device=DEV).manual_seed(seed) + lengths = spec["lengths"] + n_tok = sum(lengths) + layout = spec.get("layout", "in_proj") + + def randn(*shape): + return torch.randn(*shape, generator=g, device=DEV).to(DTYPE) + + if layout == "contiguous": + x, a, b = randn(n_tok, QKV_DIM), randn(n_tok, NUM_V_HEADS), randn(n_tok, NUM_V_HEADS) + else: + proj = randn(n_tok + spec.get("extra_rows", 0), IN_PROJ_WIDTH) + x = proj[:n_tok, :QKV_DIM] + a = proj[:n_tok, A_COL : A_COL + NUM_V_HEADS] + b = proj[:n_tok, A_COL + NUM_V_HEADS :] + assert x.stride(0) == a.stride(0) == IN_PROJ_WIDTH + weight = (torch.randn(QKV_DIM, CONV_WIDTH, generator=g, device=DEV) * 0.5).to(DTYPE) + pool_slots = spec["pool_slots"] + # The pool is a per-layer view of a two-layer buffer, as a cache manager + # hands it out: slot stride twice the slot's own size. + conv_states = randn(pool_slots, 2, QKV_DIM, CONV_WIDTH - 1)[:, 1] + slots = spec.get("slots") or _engine_slots(len(lengths)) + bounds = [0] + for n in lengths: + bounds.append(bounds[-1] + n) + qsl = torch.tensor(bounds, dtype=torch.int32, device=DEV) + si = torch.tensor(slots, dtype=torch.int32, device=DEV) + hi = torch.tensor(spec["initial"], dtype=torch.bool, device=DEV) + lo, hi_ = A_LOG_RANGE + A_log = torch.rand(NUM_V_HEADS, generator=g, device=DEV) * (hi_ - lo) + lo + lo, hi_ = DT_BIAS_RANGE + dt_bias = torch.rand(NUM_V_HEADS, generator=g, device=DEV) * (hi_ - lo) + lo + if "dt_bias_fill" in spec: + dt_bias = torch.full((NUM_V_HEADS,), spec["dt_bias_fill"], device=DEV) + return [ + x, + weight, + conv_states, + qsl, + si, + hi, + a, + b, + A_log, + dt_bias, + *HEADS, + PAD, + L2_NORM_EPS, + SOFTPLUS_THRESHOLD, + ] + + +def _unfused(args): + """The target's previous three calls on a copy of the pool: `(q, k, v, g, + beta, pool)`.""" + x, weight, conv_states, qsl, si, hi, a, b, A_log, dt_bias = args[:10] + pool = conv_states.clone() + xt = extract_transpose_prefill_slice(x, x.shape[0], 0, QKV_DIM) + causal_conv1d_fwd(xt, weight, None, pool, qsl, si, hi, True, PAD, None) + out = fused_gdn_post_conv( + xt, None, a, b, A_log, dt_bias, *HEADS, L2_NORM_EPS, SOFTPLUS_THRESHOLD + ) + return (*out, pool) + + +def _live_rows(args): + """Token rows of the contexts whose slot is not the pad sentinel.""" + bounds, slots = args[3].tolist(), args[4].tolist() + keep = torch.zeros(args[0].shape[0], dtype=torch.bool, device=DEV) + for s, slot in enumerate(slots): + if slot != PAD: + keep[bounds[s] : bounds[s + 1]] = True + return keep + + +def _compare(out, pool, want, rows=None, what=""): + q, k, v, g, beta = out + q_r, k_r, v_r, g_r, beta_r, pool_r = want + for got, ref in zip(out, want[:5]): + assert got.dtype == ref.dtype and got.shape == ref.shape and got.is_contiguous() + sel = (lambda t: t[0, rows]) if rows is not None else (lambda t: t) + for name, got, ref in (("q", q, q_r), ("k", k, k_r), ("v", v, v_r)): + d = (sel(got).float() - sel(ref).float()).abs() + flips = (d > 0).float().mean().item() if d.numel() else 0.0 + print( + f"{what} {name}: max abs {d.max().item() if d.numel() else 0:.3e}, " + f"differing fraction {flips:.2e}" + ) + torch.testing.assert_close(sel(got), sel(ref)) + torch.testing.assert_close(g, g_r, rtol=8 * _APPROX_ULP, atol=G_ATOL) + torch.testing.assert_close(beta, beta_r) + assert torch.equal(pool, pool_r), "conv pool differs" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) + pool0 = args[2].clone() + want = op.reference(*args) + old = _unfused([*args[:2], pool0.clone(), *args[3:]]) + inputs = [t.clone() if isinstance(t, torch.Tensor) else t for t in args] + out = op(*args) + for i, (arg, before) in enumerate(zip(args, inputs)): + if isinstance(arg, torch.Tensor) and i != 2: + assert torch.equal(arg, before), f"input {i} must be read-only" + _compare(out, args[2], want, what=f"{cell.why[:30]} vs reference") + _compare(out, args[2], old, rows=_live_rows(args), what=f"{cell.why[:30]} vs unfused") + if PAD in args[4].tolist(): + pad_free = [s for s in range(pool0.shape[0]) if s not in args[4].tolist()] + assert torch.equal(args[2][pad_free], pool0[pad_free]), "unnamed slots written" + + +def test_carry_and_silu_are_load_bearing() -> None: + """The band separates a dropped carry and a dropped SiLU.""" + spec = dict(lengths=(5, 70), pool_slots=4, slots=(1, 3), initial=(True, True)) + args = _build(spec, seed=301) + want = op.reference(*args) + no_carry = list(args) + no_carry[5] = torch.zeros_like(args[5]) + out = op(*[t.clone() if isinstance(t, torch.Tensor) else t for t in no_carry]) + with pytest.raises(AssertionError): + torch.testing.assert_close(out[2], want[2]) + x, weight = args[0], args[1] + # No SiLU: the pre-activation of the second context's last token, v channels. + pre = sum(x[74 - 3 + j, :QKV_DIM].float() * weight[:, j].float() for j in range(4)) + pre = pre[2 * NUM_K_HEADS * HEAD_K_DIM :] + got = op(*_build(spec, seed=301))[2][0, 74].reshape(-1).float() + assert (got - torch.nn.functional.silu(pre)).abs().max() < 0.05 + assert (got - pre).abs().max() > 0.5 + + +def test_guard_refuses_channel_strided_x() -> None: + args = _build(dict(lengths=(16,), pool_slots=2, slots=(0,), initial=(True,)), seed=302) + bad = list(args) + bad[0] = torch.empty(QKV_DIM, 16, dtype=DTYPE, device=DEV).t() + with pytest.raises(AssertionError, match="unit channel stride"): + with validating(op): + op(*bad) + + +def test_guard_refuses_short_a_b_and_strided_gates() -> None: + args = _build(dict(lengths=(16, 8), pool_slots=2, slots=(0, 1), initial=(True, True)), seed=303) + for position in (6, 7): + bad = list(args) + bad[position] = args[position][:16] + with pytest.raises(AssertionError, match="at least T=24 rows"): + with validating(op): + op(*bad) + bad = list(args) + bad[8] = torch.stack([args[8], args[8]], dim=1)[:, 0] + with pytest.raises(AssertionError, match="contiguous 1-D"): + with validating(op): + op(*bad) From d7b77d588da6fc5674e8fb6670c6f0a367c718e7 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Sun, 18 Oct 2026 23:00:02 -0700 Subject: [PATCH 07/28] [TRTLLM-16460][perf] Qwen3.8-Flash-Next target: hyper-connection ops The target's hyper-connection (HC) site kernels, with their unit tests: - Decode sites: qwen4_exp_hc_decode_tc_first (the first projection as one tcgen05 split-K GEMM at 2..128 rows), qwen4_exp_hc_decode_first / _first_cute (one-row batches), and the second halves qwen4_exp_hc_decode_mix / _mix_fp4 / _mix_fp4_cute (grouped RMS, SiLU bottleneck, injection weights, up projection; the MoE site's NVFP4 output). - qwen4_exp_hc_up_gate_mean / _up_gate_mean_fp4: the mixer's second half as one tcgen05 GEMM with the sigmoid gate and stream mean in its epilogue. - qwen4_exp_hc_inject_norm: the post-block injection fused with the next site's grouped norm; qwen4_exp_hc_lowrank_act and qwen4_exp_hc_gate_mean: the mixer's activations and epilogue. Folds these commits of the review branch (their changes to the files of this commit): - qwen3.8 target: carry the decode-fusion perf work onto the feature target (62151b727b) - qwen3.8 target: decode MoE-site inject-norm in CuTe DSL (13ea57b73e) Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../ops/hc/qwen4_exp_hc_decode_first.py | 491 ++++++++ .../ops/hc/qwen4_exp_hc_decode_first_cute.py | 525 ++++++++ .../ops/hc/qwen4_exp_hc_decode_mix.py | 331 +++++ .../ops/hc/qwen4_exp_hc_decode_mix_fp4.py | 431 +++++++ .../hc/qwen4_exp_hc_decode_mix_fp4_cute.py | 443 +++++++ .../ops/hc/qwen4_exp_hc_decode_tc_first.py | 670 ++++++++++ .../ops/hc/qwen4_exp_hc_gate_mean.py | 125 ++ .../ops/hc/qwen4_exp_hc_inject_norm.py | 1088 +++++++++++++++++ .../ops/hc/qwen4_exp_hc_lowrank_act.py | 153 +++ .../ops/hc/qwen4_exp_hc_up_gate_mean.py | 908 ++++++++++++++ .../ops/hc/qwen4_exp_hc_up_gate_mean_fp4.py | 124 ++ ...t_modeling_v2_qwen4_exp_hc_decode_first.py | 450 +++++++ ...eling_v2_qwen4_exp_hc_decode_first_cute.py | 241 ++++ ...est_modeling_v2_qwen4_exp_hc_decode_mix.py | 157 +++ ...odeling_v2_qwen4_exp_hc_decode_mix_cute.py | 132 ++ ...modeling_v2_qwen4_exp_hc_decode_mix_fp4.py | 190 +++ ...ing_v2_qwen4_exp_hc_decode_mix_fp4_cute.py | 157 +++ ...odeling_v2_qwen4_exp_hc_decode_tc_first.py | 186 +++ ...test_modeling_v2_qwen4_exp_hc_gate_mean.py | 145 +++ ...st_modeling_v2_qwen4_exp_hc_inject_norm.py | 422 +++++++ ...st_modeling_v2_qwen4_exp_hc_lowrank_act.py | 100 ++ ...t_modeling_v2_qwen4_exp_hc_up_gate_mean.py | 150 +++ ...deling_v2_qwen4_exp_hc_up_gate_mean_fp4.py | 183 +++ 23 files changed, 7802 insertions(+) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first_cute.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix_fp4.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix_fp4_cute.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_tc_first.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_gate_mean.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_inject_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_lowrank_act.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean_fp4.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_first.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_cute.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_gate_mean.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_inject_norm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_lowrank_act.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first.py new file mode 100644 index 000000000000..a307fc61006f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first.py @@ -0,0 +1,491 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Decode-row hyper-connection site, first half: the previous block's gated +injection and the mixer's first projection (`[down; block_inject]`) as one +Triton GEMV launch that leaves per-stream partial logits and per-stream sums +of squares for `qwen4_exp_hc_decode_mix`. + +Tier 5 (Triton written for this target). At decode row counts the first +projection is a GEMV over a 6.6 MB weight; cuBLAS reaches ~1.5 TB/s on it, and +the injection, the norm and the activations were three more launches. The RMS +is a per-(row, stream) scalar, so the projection runs on the *un-normed* +stream scaled by `(1 + norm_weight)`, one stream per program column: each +program's dot is exactly one stream's partial, and the consumer applies +`rstd` when it combines the four. No split-K workspace, atomics or counters. +""" + +from __future__ import annotations + +import functools +import os +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Largest row count this entry is certified for (the forward's dispatch threshold +#: must not exceed it). +MAX_ROWS = 16 + + +@functools.lru_cache(maxsize=None) +def pdl_enabled(device_index: int) -> bool: + """Programmatic dependent launch: sm_90+ unless TRTLLM_ENABLE_PDL=0.""" + props = torch.cuda.get_device_properties(device_index) + return props.major >= 9 and os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" + + +@triton.jit +def _hc_decode_first_kernel( + hc_ptr, + block_ptr, + inj_ptr, + shared_ptr, + gate_ptr, + topk_w_ptr, + topk_row_ptr, + norm_w_ptr, + w1_ptr, + hc_out_ptr, + part_ptr, + ss_ptr, + stride_hc, + stride_block, + stride_inj, + stride_shared, + stride_gate, + stride_topk_w, + stride_topk_row, + stride_hc_out, + stride_w1, + M: tl.constexpr, + HIDDEN: tl.constexpr, + N_REAL: tl.constexpr, + N_PAD: tl.constexpr, + BLOCK_N: tl.constexpr, + KC: tl.constexpr, + NC: tl.constexpr, + NC_P2: tl.constexpr, + HAS_INJECT: tl.constexpr, + HAS_SHARED: tl.constexpr, + TOPK: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, +): + # grid = (BLOCK_N-row tiles of the weight, streams). One program: BLOCK_N weight + # rows across one whole stream, as a single [BLOCK_N, NC_P2, KC] load issued + # before any arithmetic. Reductions run along the tile's last two axes, which a + # warp holds, so no CTA barrier orders a memory round trip behind another. + pid_n = tl.program_id(0) + s = tl.program_id(1) + rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + jc = tl.arange(0, NC_P2) + kc = tl.arange(0, KC) + koff = jc[:, None] * KC + kc[None, :] + kmask = (jc < NC)[:, None] & (kc >= 0)[None, :] + kcol = s * HIDDEN + koff + w = tl.load( + w1_ptr + rn[:, None, None] * stride_w1 + kcol[None, :, :], + mask=(rn < N_REAL)[:, None, None] & kmask[None, :, :], + other=0.0, + ) + gam = 1.0 + tl.load(norm_w_ptr + kcol, mask=kmask, other=0.0).to(tl.float32) + if LAUNCH_WITH_PDL: + # The weights are static, so they stream while the previous kernel finishes; + # every read of an activation and every write comes after the wait. + tl.extra.cuda.gdc_launch_dependents() + tl.extra.cuda.gdc_wait() + lead = (rn == pid_n * BLOCK_N) & (pid_n == 0) + for m in tl.static_range(M): + x = tl.load(hc_ptr + m * stride_hc + kcol, mask=kmask, other=0.0).to(tl.float32) + if HAS_INJECT: + if TOPK > 0: + # The routed experts' finalize: this row's TOPK permuted FC2 rows weighted + # by their routing weights, fp32 FMAs in slot order from zero (a slot mapped + # to -1 is skipped), stored once in the rows' dtype. + b = tl.zeros([NC_P2, KC], dtype=tl.float32) + for k in tl.static_range(TOPK): + row = tl.load(topk_row_ptr + m * stride_topk_row + k) + wk = tl.load(topk_w_ptr + m * stride_topk_w + k).to(tl.float32) + r = tl.load( + block_ptr + tl.maximum(row, 0).to(tl.int64) * stride_block + koff, + mask=kmask, + other=0.0, + ).to(tl.float32) + b = tl.where(row >= 0, tl.fma(r, wk, b), b) + b = b.to(block_ptr.dtype.element_ty).to(tl.float32) + else: + b = tl.load(block_ptr + m * stride_block + koff, mask=kmask, other=0.0).to( + tl.float32 + ) + if HAS_SHARED: + # The MoE block output, `routed + sigmoid(gate) * shared` in fp32, rounded + # once to the routed output's dtype -- the value that block used to store. + sg = tl.sigmoid(tl.load(gate_ptr + m * stride_gate).to(tl.float32)) + sh = tl.load(shared_ptr + m * stride_shared + koff, mask=kmask, other=0.0) + b = (b + sg * sh.to(tl.float32)).to(block_ptr.dtype.element_ty).to(tl.float32) + g = tl.load(inj_ptr + m * stride_inj + s).to(tl.float32) + # Rounded once, like `addcmul`: the norm and every later site read the rounded value. + xr = (x + b * g).to(hc_out_ptr.dtype.element_ty) + tl.store(hc_out_ptr + m * stride_hc_out + kcol, xr, mask=kmask & (pid_n == 0)) + x = xr.to(tl.float32) + p = tl.sum(tl.sum(w.to(tl.float32) * (x * gam)[None, :, :], axis=2), axis=1) + tl.store(part_ptr + (s * M + m) * N_PAD + rn, p) + sq = tl.sum( + tl.sum(tl.broadcast_to((x * x)[None, :, :], [BLOCK_N, NC_P2, KC]), axis=2), axis=1 + ) + tl.store(ss_ptr + s * M + m + 0 * rn, sq, mask=lead) + + +#: Launch configuration: (BLOCK_N, KC, num_warps). Swept on GB300 (see the note). +DEFAULT_CONFIG = (2, 512, 2) +#: With the MoE combine every program also reads the routed and shared rows, so +#: wider weight tiles (half the programs) win: 0.2-0.3 us per site on GB300. +COMBINE_CONFIG = (4, 512, 2) +#: With the routed experts' finalize folded in (`topk_rows`), every program +#: gathers the row's top-k FC2 rows as well. +FOLD_CONFIG = (4, 512, 2) + + +def launch( + hyper_input, + block_out, + inject, + norm_weight, + first, + hc_count, + n_cols, + shared=None, + gate_logits=None, + topk_weights=None, + topk_rows=None, + config=DEFAULT_CONFIG, +): + """The launch behind `Qwen4ExpHcDecodeFirst.__call__`, tile config exposed. + Returns `(hc_out, part, ss)`.""" + rows, width = hyper_input.shape + hidden = width // hc_count + block_n, kc, warps = config + while hidden % kc: + kc //= 2 + n_tiles = triton.cdiv(n_cols, block_n) + n_pad = n_tiles * block_n + dev, dt = hyper_input.device, hyper_input.dtype + has_inject = block_out is not None + has_shared = shared is not None + top_k = topk_rows.shape[1] if topk_rows is not None else 0 + hc_out = torch.empty((rows, width), dtype=dt, device=dev) if has_inject else hyper_input + part = torch.empty((hc_count, rows, n_pad), dtype=torch.float32, device=dev) + ss = torch.empty((hc_count, rows), dtype=torch.float32, device=dev) + pdl = pdl_enabled(dev.index if dev.index is not None else torch.cuda.current_device()) + if rows > 0: + _hc_decode_first_kernel[(n_tiles, hc_count)]( + hyper_input, + block_out if has_inject else hyper_input, + inject if has_inject else hyper_input, + shared if has_shared else hyper_input, + gate_logits if has_shared else hyper_input, + topk_weights if top_k else hyper_input, + topk_rows if top_k else hyper_input, + norm_weight, + first, + hc_out, + part, + ss, + hyper_input.stride(0), + block_out.stride(0) if has_inject else 0, + inject.stride(0) if has_inject else 0, + shared.stride(0) if has_shared else 0, + gate_logits.stride(0) if has_shared else 0, + topk_weights.stride(0) if top_k else 0, + topk_rows.stride(0) if top_k else 0, + hc_out.stride(0), + first.stride(1), + M=rows, + HIDDEN=hidden, + N_REAL=n_cols, + N_PAD=n_pad, + BLOCK_N=block_n, + KC=kc, + NC=hidden // kc, + NC_P2=triton.next_power_of_2(hidden // kc), + HAS_INJECT=has_inject, + HAS_SHARED=has_shared, + TOPK=top_k, + LAUNCH_WITH_PDL=pdl, + num_warps=warps, + launch_pdl=pdl, + ) + return hc_out, part, ss + + +class Qwen4ExpHcDecodeFirst(OpWrapper): + """With `C = hc_count`, `H = width // C`, `N = n_cols`, `first` the + `[width, >= N]` view of the load-time `[down; block_inject; 0]` weight:: + + hc_out[t, s*H+i] = hyper_input[t, s*H+i] + block_out[t, i] * inject[t, s] (rounded once) + part[s, t, n] = sum_i hc_out[t, s*H+i] * (1 + norm_weight[s*H+i]) * first[s*H+i, n] + ss[s, t] = sum_i hc_out[t, s*H+i]^2 + + Returns `(hc_out, part, ss)`: `hc_out` fresh in the activation dtype (it + *is* `hyper_input` when `block_out` / `inject` are None -- the norm-only + site); `part` fp32 `[C, T, N]` (a column slice of a padded buffer) and `ss` + fp32 `[C, T]`. The grouped RMS `rstd[t, s] = (ss[s, t] / H + eps)^-1/2` + scales stream `s`'s partial: `normed @ first[:, :N] = sum_s rstd[:, s] * + part[s]`. `qwen4_exp_hc_decode_mix` makes that combination. + + With `shared` and `gate_logits` (the site after a MoE block), `block_out` + is the routed experts' output and the block output injected above is + `block_out[t, i] + sigmoid(gate_logits[t]) * shared[t, i]` in fp32, rounded + once to `block_out`'s dtype -- bitwise what `fused_sigmoid_gate_mul_add` + stores. `gate_logits` is one pre-sigmoid scalar per row, `[T, 1]` or `[T]`. + + With `topk_weights` / `topk_rows` as well (the MoE runner left unfinalized), + `block_out` is its permuted FC2 output `[P, H]` and the routed row is the + runner's finalize, formed here: `routed[t, i] = sum_k topk_weights[t, k] * + block_out[topk_rows[t, k], i]` as fp32 FMAs in slot order from zero (a slot + whose row is -1 is skipped), rounded once to `block_out`'s dtype -- bitwise + what the runner's finalize stores. `topk_rows` is the runner's + expanded->permuted index map `[T, k]` int32, `topk_weights` its `[T, k]` + routing weights. The fold needs the combine (`shared`, `gate_logits`). + + Fusion boundary: the routed experts' finalize (optional), the + shared-expert combine (optional), the previous block's + injection, the `(1 + w)` scale and the first projection per stream, plus + the per-stream sum of squares. The routed and shared experts and the gate + projection, the RMS itself, the activations and everything after are not + inside. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: rows, inject (True), moe (False: shared + gate_logits too), gate + # "column" [T, 1] | "flat" [T], fold (0: top-k of the finalize folded in, + # block_out then the permuted rows), drop (False: some slots -1), n_cols + # (324), hc_count (4), hidden (2560), n_rows (384: the padded weight), + # dtype (bf16). + CELLS: tuple[Cell, ...] = tuple( + Cell(why=f"checkpoint 4x2560, 324 cols of 384, bf16, inject, {why}", spec=dict(rows=t)) + for t, why in ( + (1, "decode 1 row (con 1)"), + (2, "decode 2 rows"), + (4, "decode 4 rows"), + (8, "decode 8 rows"), + (16, "decode 16 rows (MAX_ROWS)"), + ) + ) + ( + Cell(why="MoE combine, [T, 1] gate, 1 row (con 1)", spec=dict(rows=1, moe=True)), + Cell(why="MoE combine, flat [T] gate, 4 rows", spec=dict(rows=4, moe=True, gate="flat")), + Cell(why="MoE combine, 16 rows (MAX_ROWS)", spec=dict(rows=16, moe=True)), + Cell( + why="MoE combine into the stack-final mixer: 320 cols, 1 row", + spec=dict(rows=1, moe=True, n_cols=320), + ), + Cell(why="finalize fold, top-10, 1 row (con 1)", spec=dict(rows=1, moe=True, fold=10)), + Cell( + why="finalize fold into the stack-final mixer: 320 cols, 1 row", + spec=dict(rows=1, moe=True, fold=10, n_cols=320), + ), + Cell( + why="finalize fold, 4 rows, slots mapped to -1, flat gate", + spec=dict(rows=4, moe=True, fold=10, drop=True, gate="flat"), + ), + Cell(why="finalize fold, 16 rows (MAX_ROWS)", spec=dict(rows=16, moe=True, fold=10)), + Cell( + why="norm-only site (stack start / after PLE), 1 row", spec=dict(rows=1, inject=False) + ), + Cell(why="norm-only site, 8 rows", spec=dict(rows=8, inject=False)), + Cell(why="stack-final mixer: 320 cols, 1 row", spec=dict(rows=1, n_cols=320)), + Cell(why="stack-final mixer: 320 cols, 16 rows", spec=dict(rows=16, n_cols=320)), + Cell(why="hyper_input a row slice buf[::2], 4 rows", spec=dict(rows=4, layout="row_slice")), + Cell(why="unpadded weight [width, 324], 3 rows", spec=dict(rows=3, n_rows=324)), + Cell(why="hc_count 2 at hidden 2560, 2 rows", spec=dict(rows=2, hc_count=2)), + Cell( + why="hidden 128, 36 cols (not the checkpoint), 5 rows", + spec=dict(rows=5, hidden=128, n_cols=36), + ), + Cell(why="float16 operands, 4 rows", spec=dict(rows=4, dtype=torch.float16)), + Cell(why="zero rows", spec=dict(rows=0)), + ) + + note = """ + Metadata consumed: none. Call sites: `DecodeTarget`'s trunk, every + hyper-connection site (97 per forward) when the decode batch has at most + `_HC_DECODE_GEMV_MAX_ROWS` rows (<= MAX_ROWS = 16), followed by + `qwen4_exp_hc_decode_mix`. Larger batches and `PrefillTarget` keep + `qwen4_exp_hc_inject_norm` + `cublas_mm` + `qwen4_exp_hc_lowrank_act`. + The 47 sites after a MoE block other than the one before the PLE add pass + `shared` and `gate_logits`, and with them the unfinalized runner's + `topk_weights` / `topk_rows` (`block_out` its permuted FC2 rows). + + Preconditions: 2-D operands with unit column stride (row strides free); + `first` `[width, >= n_cols]` with unit stride along `width` (the `.t()` + view of a row-major `[N, width]` weight); `hidden` a multiple of 16. Row + count is a compile-time constant (one Triton specialization per decode + batch size); at most MAX_ROWS rows certified. Columns past `n_cols` are + never read, so the zero padding of the load-time weight costs no traffic. + + Numerics: `hc_out` as `qwen4_exp_hc_inject_norm` (fp32 accumulate, one + rounding); the projection multiplies the fp32 `hc_out * (1 + w)` by the + 16-bit weight in fp32 (the composition it replaces rounded the normed + stream to bf16 first); deterministic (no atomics, fixed reduction order). + The optional combine is fp32 rounded once to `block_out`'s dtype, which + against the float64 reference can land one ulp of the block output away + (its rounding midpoint); `hc_out` carries that scaled by `inject`, and the + test checks `part` / `ss` against the projection of the `hc_out` the op + stored. Every program recomputes the combine for its stream's columns: the + extra `[hidden]` read per program is L2-resident. The optional finalize is + the runner's own arithmetic (fp32 FMAs in slot order from zero, one + rounding), so `hc_out` and `ss` are bitwise the finalize + combine + composition at the same tile, and so is `part` from 2 rows; at 1 row the + fold's Triton specialization reduces the dot in another fp32 order (~12% of + `part` one fp32 ulp apart, < 1e-7 relative; GB300). The + float64 reference rounds the finalize from an exact sum, which can land one + ulp of the routed row away, and the test's band carries that term too. + With the fold every program also gathers the row's top-k FC2 rows (`k x + [hidden]`, L2-resident: the runner's GEMM wrote them just before). + Graph-capturable: no host sync, outputs allocated per call. Launched with + programmatic dependent launch on sm_90+ (`TRTLLM_ENABLE_PDL=0` turns it + off): the static weight tile streams before `gdc_wait`, every activation + read and every write after it, and dependents are released at once so + `qwen4_exp_hc_decode_mix` can stream its weights meanwhile. + + Grid: (cdiv(n_cols, 2), hc_count) programs of 2 warps, each one load of 2 + weight rows x one stream (10 KB); with the combine, 4 weight rows + (`COMBINE_CONFIG`; it sums `part` in a different fp32 order than the + default tile, inside the same band). Measured on GB300 (graph replay, cold + weights, with `qwen4_exp_hc_decode_mix`): 1 row is fast; each further row + costs several us (the row loop serializes its memory round trips), so the + forward dispatches only 1-row batches here. Hard limits: int32 offsets, far below 2^31 at <= 16 rows. + """ + + def __call__( + self, + hyper_input: torch.Tensor, + block_out: Optional[torch.Tensor], + inject: Optional[torch.Tensor], + norm_weight: torch.Tensor, + first: torch.Tensor, + hc_count: int, + n_cols: int, + shared: Optional[torch.Tensor] = None, + gate_logits: Optional[torch.Tensor] = None, + topk_weights: Optional[torch.Tensor] = None, + topk_rows: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if topk_rows is not None: + config = FOLD_CONFIG + elif shared is not None: + config = COMBINE_CONFIG + else: + config = DEFAULT_CONFIG + hc_out, part, ss = launch( + hyper_input, + block_out, + inject, + norm_weight, + first, + hc_count, + n_cols, + shared, + gate_logits, + topk_weights, + topk_rows, + config=config, + ) + return hc_out, part[:, :, :n_cols], ss + + def reference( + self, + hyper_input: torch.Tensor, + block_out: Optional[torch.Tensor], + inject: Optional[torch.Tensor], + norm_weight: torch.Tensor, + first: torch.Tensor, + hc_count: int, + n_cols: int, + shared: Optional[torch.Tensor] = None, + gate_logits: Optional[torch.Tensor] = None, + topk_weights: Optional[torch.Tensor] = None, + topk_rows: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """float64; the finalize (with `topk_rows`) rounded once to + `block_out`'s dtype, the combine (with `shared`) rounded once to `block_out`'s + dtype, `hc_out` rounded once to the activation dtype, `part` and `ss` + rounded once to fp32.""" + dtype = hyper_input.dtype + rows, width = hyper_input.shape + hidden = width // hc_count + if block_out is None: + hc_out = hyper_input + else: + if topk_rows is not None: + valid = (topk_rows >= 0).unsqueeze(-1) + rows_ = block_out.double()[topk_rows.clamp(min=0).long()] + w = topk_weights.double().unsqueeze(-1) + block_out = torch.where(valid, rows_ * w, 0.0).sum(1).to(block_out.dtype) + if shared is not None: + sg = torch.sigmoid(gate_logits.double().reshape(rows, 1)) + block_out = (block_out.double() + sg * shared.double()).to(block_out.dtype) + streams = hyper_input.double().unflatten(-1, (hc_count, hidden)) + gated = block_out.double().unsqueeze(-2) * inject.double().unsqueeze(-1) + hc_out = (streams + gated).flatten(-2).to(dtype) + x = hc_out.double().unflatten(-1, (hc_count, hidden)) + y = x * (1.0 + norm_weight.double()).unflatten(-1, (hc_count, hidden)) + w = first[:, :n_cols].double().unflatten(0, (hc_count, hidden)) + part = torch.einsum("tsi,sin->stn", y, w).float() + ss = x.pow(2).sum(dim=-1).t().float() + return hc_out, part, ss + + def is_valid( + self, + hyper_input: torch.Tensor, + block_out: Optional[torch.Tensor], + inject: Optional[torch.Tensor], + norm_weight: torch.Tensor, + first: torch.Tensor, + hc_count: int, + n_cols: int, + shared: Optional[torch.Tensor] = None, + gate_logits: Optional[torch.Tensor] = None, + topk_weights: Optional[torch.Tensor] = None, + topk_rows: Optional[torch.Tensor] = None, + ) -> None: + # Raw row strides with unit column stride everywhere; a contiguous + # [width, N] weight would be read with its rows and columns swapped. + assert hyper_input.dim() == 2 and hyper_input.stride(1) == 1, "hyper_input [T, width]" + rows, width = hyper_input.shape + assert width % hc_count == 0 and (width // hc_count) % 16 == 0, "hidden multiple of 16" + assert (block_out is None) == (inject is None), "block_out and inject come together" + assert (shared is None) == (gate_logits is None), "shared and gate_logits come together" + assert shared is None or block_out is not None, "the combine needs block_out" + assert (topk_weights is None) == (topk_rows is None), ( + "topk_weights and topk_rows come together" + ) + assert topk_rows is None or shared is not None, "the finalize fold needs the combine" + if block_out is not None: + if topk_rows is None: + assert block_out.shape == (rows, width // hc_count) and block_out.stride(1) == 1 + else: + # The permuted FC2 rows: any row count, read through topk_rows. + assert block_out.dim() == 2 and block_out.shape[1] == width // hc_count + assert block_out.stride(1) == 1, "permuted rows [P, hidden]" + assert topk_rows.dtype == torch.int32, "topk_rows int32 (the runner's index map)" + assert topk_rows.dim() == 2 and topk_rows.shape[0] == rows, "topk_rows [T, k]" + assert topk_rows.stride(1) == 1 and topk_weights.stride(1) == 1 + assert topk_weights.shape == topk_rows.shape, "topk_weights [T, k]" + assert inject.shape == (rows, hc_count) and inject.stride(1) == 1 + if shared is not None: + assert shared.shape == (rows, width // hc_count), "shared [T, hidden]" + assert shared.stride(1) == 1, "shared [T, hidden]" + assert shared.dtype == block_out.dtype, "shared in block_out's dtype" + assert gate_logits.shape in ((rows,), (rows, 1)), "gate_logits [T, 1] or [T]" + assert rows <= 1 or gate_logits.stride(0) != 0, "one gate logit per row" + assert norm_weight.shape == (width,) and norm_weight.is_contiguous(), "norm_weight [width]" + assert first.shape[0] == width and first.shape[1] >= n_cols, "first [width, >= n_cols]" + assert first.stride(0) == 1, "first must be the .t() view of a row-major [N, width] weight" + + +qwen4_exp_hc_decode_first = Qwen4ExpHcDecodeFirst() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first_cute.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first_cute.py new file mode 100644 index 000000000000..8d28cf523702 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first_cute.py @@ -0,0 +1,525 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Decode-row hyper-connection site, first half, for one-row batches: the previous +block's injection (with the MoE block's combine and the routed experts' finalize +folded in, when given) and the first projection per stream, one PDL launch. + +Tier 5 (CuTe DSL written for this target). `qwen4_exp_hc_decode_first` computes the +same thing for up to 16 rows in Triton, whose power-of-two tiles pad a +2560-wide stream to 8 x 512: three eighths of every program's weight, activation +and scale tiles are masked lanes holding registers. Here a stream is exactly +`THREADS x VEC` elements, one 16-byte vector per thread, and the grid is sized so +that every CTA is resident at once (one wave), so all of them stream their weight +before the producer finishes. + +Order of work inside a CTA (`OUTS` weight rows x one stream): + +1. its weight rows and its slice of the norm weight into registers -- static, so + before `griddepcontrol.wait`, while the kernel ahead of this one runs; +2. `griddepcontrol.wait`; then the stream's vector of the residual and of the + block output -- with the fold, the slot map and the weights first, then the + slot rows -- and the injection weight; +3. the injected residual rounded once (stored by the stream's first CTA), its sum + of squares and its `(1 + w)`-scaled products with the weight rows, in fp32; +4. a warp then a cross-warp reduction; thread `j < OUTS` stores partial `j`, thread + `OUTS` of the stream's first CTA the sum of squares. +""" + +from __future__ import annotations + +import functools +import os +from typing import Optional, Tuple + +import cuda.bindings.driver as cuda_drv +import cutlass +import cutlass.cute as cute +import torch +from cutlass import const_expr +from cutlass.cute import experimental as cute_ext +from cutlass.cute.runtime import from_dlpack + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .qwen4_exp_hc_decode_first import qwen4_exp_hc_decode_first as _triton_first + +#: One decode row per launch (the 1-row batch is the one the Triton entry is slowest at +#: relative to its weight stream). +MAX_ROWS = 1 +VEC = 8 +#: Threads per CTA: one 16-byte vector of a 2560-wide stream each. +THREADS = 320 +#: Streaming multiprocessors the one-wave grid must fit (GB300). +_SMS = 152 +_COMPILE_OPTIONS = "--ptxas-options -maxrregcount=128" + + +@functools.lru_cache(maxsize=None) +def pdl_enabled(device_index: int) -> bool: + """Programmatic dependent launch: sm_90+ unless TRTLLM_ENABLE_PDL=0.""" + props = torch.cuda.get_device_properties(device_index) + return props.major >= 9 and os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" + + +def outs_for(n_cols: int, hc_count: int) -> int: + """Weight rows per CTA: the fewest that divide `n_cols` and keep the grid in one wave.""" + for outs in range(1, n_cols + 1): + if n_cols % outs == 0 and (n_cols // outs) * hc_count <= _SMS: + return outs + raise AssertionError(f"no one-wave tiling for {n_cols} columns") + + +class _FirstKernel: + def __init__(self, *, hidden, outs, inject, combine, topk, pdl): + assert hidden == THREADS * VEC, "one vector per thread" + self.hidden, self.outs = hidden, outs + self.inject, self.combine, self.topk, self.pdl = inject, combine, topk, pdl + self.warps = THREADS // 32 + + @cute.jit + def __call__( + self, + mHc, + mBlock, + mInj, + mShared, + mGate, + mTopW, + mTopRows, + mNorm, + mW, + mHcOut, + mPart, + mSs, + stream: cuda_drv.CUstream, + ): + n_tiles = cute.size(mPart, mode=[1]) // self.outs + streams = cute.size(mPart, mode=[0]) + load_static = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + cutlass.BFloat16, + num_bits_per_copy=16 * VEC, + load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING, + ) + load_act = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + cutlass.BFloat16, + num_bits_per_copy=16 * VEC, + load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS, + ) + store = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.BFloat16, num_bits_per_copy=16 * VEC + ) + self.kernel( + mHc, + mBlock, + mInj, + mShared, + mGate, + mTopW, + mTopRows, + mNorm, + mW, + mHcOut, + mPart, + mSs, + load_static, + load_act, + store, + ).launch( + grid=[n_tiles, streams, 1], + block=[THREADS, 1, 1], + smem=(self.outs + 1) * self.warps * 4, + stream=stream, + use_pdl=self.pdl, + # At most 96 registers a thread: the mix kernel that follows (512 threads at + # 56) then fits on the SM beside this CTA and streams its weights before + # this kernel ends. + min_blocks_per_mp=2, + ) + + @cute.kernel + def kernel( + self, + mHc, + mBlock, + mInj, + mShared, + mGate, + mTopW, + mTopRows, + mNorm, + mW, + mHcOut, + mPart, + mSs, + load_static: cute.CopyAtom, + load_act: cute.CopyAtom, + store: cute.CopyAtom, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, s, _ = cute.arch.block_idx() + warp = cute.arch.warp_idx() + outs: cutlass.Constexpr = self.outs + warps: cutlass.Constexpr = self.warps + topk: cutlass.Constexpr = self.topk + n0 = bidx * outs + vi = s * THREADS + tidx # this thread's vector of the width; `tidx` of the stream + + w_vec = cute.logical_divide(mW, (None, VEC)) + hc_vec = cute.logical_divide(mHc, (None, VEC)) + nw_vec = cute.logical_divide(mNorm, (None, VEC)) + + # 1. Static operands, before the wait. + w_reg = cute.make_rmem_tensor( + cute.make_layout((outs, VEC), stride=(VEC, 1)), cutlass.BFloat16 + ) + for j in cutlass.range_constexpr(outs): + cute.copy(load_static, w_vec[n0 + j, (None, vi)], w_reg[j, None]) + nw_reg = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.BFloat16) + cute.copy(load_static, nw_vec[0, (None, vi)], nw_reg) + if const_expr(self.pdl): + cute.arch.griddepcontrol_launch_dependents() + cute.arch.griddepcontrol_wait() + + # 2. The producers' outputs. + x_reg = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.BFloat16) + cute.copy(load_act, hc_vec[0, (None, vi)], x_reg) + xr = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.Float32) + for v in cutlass.range_constexpr(VEC): + xr[v] = x_reg[v].to(cutlass.Float32) + if const_expr(self.inject): + b_vec = cute.logical_divide(mBlock, (None, VEC)) + g = mInj[0, s].to(cutlass.Float32) + b = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.Float32) + if const_expr(topk > 0): + # The runner's finalize: fp32 FMAs in slot order from zero, a slot mapped + # to -1 skipped, one rounding. Every slot row is requested before any FMA. + rows = cute.make_rmem_tensor(cute.make_layout((topk,), stride=(1,)), cutlass.Int32) + wts = cute.make_rmem_tensor(cute.make_layout((topk,), stride=(1,)), cutlass.Float32) + for k in cutlass.range_constexpr(topk): + rows[k] = mTopRows[0, k] + wts[k] = mTopW[0, k].to(cutlass.Float32) + fc2 = cute.make_rmem_tensor( + cute.make_layout((topk, VEC), stride=(VEC, 1)), cutlass.BFloat16 + ) + for k in cutlass.range_constexpr(topk): + cute.copy( + load_act, + b_vec[cutlass.max(rows[k], cutlass.Int32(0)), (None, tidx)], + fc2[k, None], + ) + for v in cutlass.range_constexpr(VEC): + acc = cutlass.Float32(0.0) + for k in cutlass.range_constexpr(topk): + if rows[k] >= 0: + acc = cute.math.fma(fc2[k, v].to(cutlass.Float32), wts[k], acc) + b[v] = acc.to(cutlass.BFloat16).to(cutlass.Float32) + else: + blk = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.BFloat16) + cute.copy(load_act, b_vec[0, (None, tidx)], blk) + for v in cutlass.range_constexpr(VEC): + b[v] = blk[v].to(cutlass.Float32) + if const_expr(self.combine): + # The MoE block output, routed + sigmoid(gate) * shared in fp32, rounded once. + sh_vec = cute.logical_divide(mShared, (None, VEC)) + sh = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.BFloat16) + cute.copy(load_act, sh_vec[0, (None, tidx)], sh) + z = mGate[0].to(cutlass.Float32) + sg = 1.0 / (1.0 + cute.math.exp(-z)) + for v in cutlass.range_constexpr(VEC): + b[v] = ( + cute.math.fma(sg, sh[v].to(cutlass.Float32), b[v]) + .to(cutlass.BFloat16) + .to(cutlass.Float32) + ) + out = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.BFloat16) + for v in cutlass.range_constexpr(VEC): + # Rounded once, like `addcmul`: the norm and every later site read this value. + out[v] = cute.math.fma(b[v], g, xr[v]).to(cutlass.BFloat16) + xr[v] = out[v].to(cutlass.Float32) + if bidx == 0: + o_vec = cute.logical_divide(mHcOut, (None, VEC)) + cute.copy(store, out, o_vec[0, (None, vi)]) + + # 3. Sum of squares and the scaled products, fp32. + acc = cute.make_rmem_tensor(cute.make_layout((outs + 1,), stride=(1,)), cutlass.Float32) + acc.fill(0.0) + for v in cutlass.range_constexpr(VEC): + y = xr[v] * (1.0 + nw_reg[v].to(cutlass.Float32)) + for j in cutlass.range_constexpr(outs): + acc[j] = y * w_reg[j, v].to(cutlass.Float32) + acc[j] + acc[outs] = xr[v] * xr[v] + acc[outs] + + # 4. Reductions; thread j < OUTS stores partial j, thread OUTS the sum of squares. + for j in cutlass.range_constexpr(outs + 1): + acc[j] = cute.arch.warp_reduction_sum(acc[j]) + smem = cutlass.utils.SmemAllocator() + red = smem.allocate_tensor( + cutlass.Float32, + cute.make_layout((outs + 1, warps), stride=(warps, 1)), + byte_alignment=16, + ) + with cute.arch.elect_one(): + for j in cutlass.range_constexpr(outs + 1): + red[j, warp] = acc[j] + cute.arch.sync_threads() + if tidx < outs: + total = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(warps): + total = total + red[tidx, w] + mPart[s, n0 + tidx] = total + if tidx == outs: + if bidx == 0: + total = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(warps): + total = total + red[outs, w] + mSs[s] = total + + +def _dl(t: torch.Tensor, align: int = 16): + return from_dlpack(t, assumed_align=align) + + +@functools.lru_cache(maxsize=None) +def _compiled( + hidden, hc_count, n_cols, n_rows, inject, combine, topk, n_perm, pdl, hc_stride, block_stride +): + """One compiled launch per variant and geometry (strides are part of a CuTe layout).""" + outs = outs_for(n_cols, hc_count) + kern = _FirstKernel( + hidden=hidden, outs=outs, inject=inject, combine=combine, topk=topk, pdl=pdl + ) + dev, bf = "cuda", torch.bfloat16 + width = hidden * hc_count + hc = torch.empty(1, hc_stride, dtype=bf, device=dev)[:, :width] + block = torch.empty(max(n_perm, 1), block_stride, dtype=bf, device=dev)[:, :hidden] + inj = torch.empty(1, hc_count, dtype=bf, device=dev) + shared = torch.empty(1, hidden, dtype=bf, device=dev) + gate = torch.empty(1, dtype=bf, device=dev) + topw = torch.empty(1, max(topk, 1), dtype=bf, device=dev) + toprows = torch.empty(1, max(topk, 1), dtype=torch.int32, device=dev) + norm = torch.empty(1, width, dtype=bf, device=dev) + w = torch.empty(n_rows, width, dtype=bf, device=dev) + hc_out = torch.empty(1, width, dtype=bf, device=dev) + part = torch.empty(hc_count, n_cols, dtype=torch.float32, device=dev) + ss = torch.empty(hc_count, dtype=torch.float32, device=dev) + stream = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream) + return cute_ext.compile( + kern, + _dl(hc), + _dl(block), + _dl(inj, 2), + _dl(shared), + _dl(gate, 2), + _dl(topw, 2), + _dl(toprows, 4), + _dl(norm), + _dl(w), + _dl(hc_out), + _dl(part, 4), + _dl(ss, 4), + stream, + options=_COMPILE_OPTIONS, + ) + + +def launch( + hyper_input, + block_out, + inject, + norm_weight, + first, + hc_count, + n_cols, + shared=None, + gate_logits=None, + topk_weights=None, + topk_rows=None, +): + """The launch behind `Qwen4ExpHcDecodeFirstCute.__call__`. Returns `(hc_out, part, ss)`, + `part` fp32 `[C, 1, n_cols]` and `ss` fp32 `[C, 1]`.""" + rows, width = hyper_input.shape + hidden = width // hc_count + dev = hyper_input.device + has_inject = block_out is not None + has_shared = shared is not None + top_k = topk_rows.shape[1] if topk_rows is not None else 0 + hc_out = ( + torch.empty((rows, width), dtype=hyper_input.dtype, device=dev) + if has_inject + else hyper_input + ) + part = torch.empty((hc_count, rows, n_cols), dtype=torch.float32, device=dev) + ss = torch.empty((hc_count, rows), dtype=torch.float32, device=dev) + if rows == 0: + return hc_out, part, ss + pdl = pdl_enabled(dev.index if dev.index is not None else torch.cuda.current_device()) + weight = first.t() # the row-major [n_rows, width] weight + gate = gate_logits.reshape(1) if has_shared else hyper_input.new_empty(1) + n_perm = block_out.shape[0] if (has_inject and top_k) else 1 + fn = _compiled( + hidden, + hc_count, + n_cols, + weight.shape[0], + has_inject, + has_shared, + top_k, + n_perm, + bool(pdl), + hyper_input.stride(0), + block_out.stride(0) if (has_inject and top_k) else hidden, + ) + blk = block_out if has_inject else hyper_input.new_empty(1, hidden) + inj = inject if has_inject else hyper_input.new_empty(1, hc_count) + sh = shared if has_shared else hyper_input.new_empty(1, hidden) + topw = topk_weights if top_k else hyper_input.new_empty(1, 1) + toprows = topk_rows if top_k else torch.empty(1, 1, dtype=torch.int32, device=dev) + stream = cuda_drv.CUstream(torch.cuda.current_stream(dev).cuda_stream) + fn( + _dl(hyper_input), + _dl(blk), + _dl(inj, 2), + _dl(sh), + _dl(gate, 2), + _dl(topw, 2), + _dl(toprows, 4), + _dl(norm_weight.reshape(1, -1)), + _dl(weight), + _dl(hc_out), + _dl(part[:, 0, :], 4), + _dl(ss[:, 0], 4), + stream, + ) + return hc_out, part, ss + + +class Qwen4ExpHcDecodeFirstCute(OpWrapper): + """`qwen4_exp_hc_decode_first`'s computation and return values for a one-row batch + (see that entry for the formulas, the MoE combine and the finalize fold):: + + hc_out[0, s*H+i] = hyper_input[0, s*H+i] + block_out[0, i] * inject[0, s] (rounded once) + part[s, 0, n] = sum_i hc_out[0, s*H+i] * (1 + norm_weight[s*H+i]) * first[s*H+i, n] + ss[s, 0] = sum_i hc_out[0, s*H+i]^2 + + Returns `(hc_out, part, ss)`: `hc_out` fresh (it *is* `hyper_input` at a norm-only + site), `part` fp32 `[C, 1, N]`, `ss` fp32 `[C, 1]`. The combine (`shared`, + `gate_logits`) and the fold (`topk_weights`, `topk_rows`) are the Triton entry's. + + Fusion boundary: as `qwen4_exp_hc_decode_first`. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: the Triton entry's (rows is 1). + CELLS: tuple[Cell, ...] = ( + Cell(why="checkpoint 4x2560, 324 cols of 384, bf16, inject, 1 row", spec=dict(rows=1)), + Cell(why="MoE combine, [T, 1] gate, 1 row", spec=dict(rows=1, moe=True)), + Cell(why="MoE combine, flat [T] gate, 1 row", spec=dict(rows=1, moe=True, gate="flat")), + Cell( + why="MoE combine into the stack-final mixer: 320 cols", + spec=dict(rows=1, moe=True, n_cols=320), + ), + Cell(why="finalize fold, top-10, 1 row (con 1)", spec=dict(rows=1, moe=True, fold=10)), + Cell( + why="finalize fold into the stack-final mixer: 320 cols", + spec=dict(rows=1, moe=True, fold=10, n_cols=320), + ), + Cell( + why="finalize fold, last 3 slots mapped to -1", + spec=dict(rows=1, moe=True, fold=10, drop=True), + ), + Cell(why="norm-only site (stack start / after PLE)", spec=dict(rows=1, inject=False)), + Cell(why="stack-final mixer: 320 cols", spec=dict(rows=1, n_cols=320)), + Cell(why="hyper_input a row slice buf[::2]", spec=dict(rows=1, layout="row_slice")), + Cell(why="unpadded weight [width, 324]", spec=dict(rows=1, n_rows=324)), + ) + + note = """ + Metadata consumed: none. Call sites: `DecodeTarget`'s trunk, every + hyper-connection site of a one-row decode step (97 per forward), followed by + `qwen4_exp_hc_decode_mix`; batches of more rows keep `qwen4_exp_hc_decode_first`. + + Preconditions: as the Triton entry, and: one row; `hidden` 2560 (320 threads x + 8); bf16 operands; `n_cols` with a divisor that keeps `n_cols / outs x C <= 152` + CTAs (324 -> 9 rows x 36 tiles x 4 streams = 144; 320 -> 10 rows, 128); the + permuted rows' count is part of the compiled layout (one compile per value). + + Numerics: `hc_out` the Triton entry's arithmetic (fp32 FMAs, one rounding; the + finalize and the combine likewise), so bitwise equal to it; `part` and `ss` + reduce in another fp32 order (inside the same band of the float64 reference). + Deterministic. Graph-capturable: compiled on the first call per variant, which + must precede capture (the engine's warmup forward does it). + + PDL: with `TRTLLM_ENABLE_PDL` (default on, sm_90+) the weight rows and the norm + slice stream before `griddepcontrol.wait`, every activation read and every + write after it, and dependents are released before the wait. + """ + + def __call__( + self, + hyper_input: torch.Tensor, + block_out: Optional[torch.Tensor], + inject: Optional[torch.Tensor], + norm_weight: torch.Tensor, + first: torch.Tensor, + hc_count: int, + n_cols: int, + shared: Optional[torch.Tensor] = None, + gate_logits: Optional[torch.Tensor] = None, + topk_weights: Optional[torch.Tensor] = None, + topk_rows: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return launch( + hyper_input, + block_out, + inject, + norm_weight, + first, + hc_count, + n_cols, + shared, + gate_logits, + topk_weights, + topk_rows, + ) + + def reference(self, *args, **kwargs) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """The Triton entry's float64 reference.""" + return _triton_first.reference(*args, **kwargs) + + def is_valid( + self, + hyper_input: torch.Tensor, + block_out: Optional[torch.Tensor], + inject: Optional[torch.Tensor], + norm_weight: torch.Tensor, + first: torch.Tensor, + hc_count: int, + n_cols: int, + shared: Optional[torch.Tensor] = None, + gate_logits: Optional[torch.Tensor] = None, + topk_weights: Optional[torch.Tensor] = None, + topk_rows: Optional[torch.Tensor] = None, + ) -> None: + _triton_first.is_valid( + hyper_input, + block_out, + inject, + norm_weight, + first, + hc_count, + n_cols, + shared, + gate_logits, + topk_weights, + topk_rows, + ) + assert hyper_input.shape[0] <= MAX_ROWS, "one row" + assert hyper_input.shape[1] // hc_count == THREADS * VEC, "hidden 2560" + assert hyper_input.dtype == torch.bfloat16, "bf16" + outs_for(n_cols, hc_count) + + +qwen4_exp_hc_decode_first_cute = Qwen4ExpHcDecodeFirstCute() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix.py new file mode 100644 index 000000000000..144ceea58438 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix.py @@ -0,0 +1,331 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Decode-row hyper-connection site, second half: the grouped RMS, the SiLU +bottleneck and the injection weights from `qwen4_exp_hc_decode_first`'s +per-stream partials, then the up projection, the sigmoid gate on the normed +stream and the mean over streams -- one Triton GEMV launch. + +Tier 5 (Triton written for this target). Each program owns `BLOCK_H` hidden +columns across all streams: it requests its whole up-projection tile before +any arithmetic, rebuilds `rstd` and `low` from the four partials (a few KB of +L2 per row, recomputed by every program), and keeps the `[T, C*H]` gate +logits in registers. Program 0 also writes the block-injection weights. +""" + +from __future__ import annotations + +import functools +import os +from typing import Optional, Tuple + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Largest row count this entry is certified for. +MAX_ROWS = 16 + + +@functools.lru_cache(maxsize=None) +def pdl_enabled(device_index: int) -> bool: + """Programmatic dependent launch: sm_90+ unless TRTLLM_ENABLE_PDL=0.""" + props = torch.cuda.get_device_properties(device_index) + return props.major >= 9 and os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" + + +@triton.jit +def _hc_decode_mix_kernel( + part_ptr, + ss_ptr, + up_ptr, + x_ptr, + norm_w_ptr, + mixed_ptr, + inj_out_ptr, + stride_ps, + stride_pm, + stride_ss, + stride_up_r, + stride_up_n, + stride_x, + stride_mixed, + eps, + M: tl.constexpr, + HIDDEN: tl.constexpr, + HC: tl.constexpr, + R1: tl.constexpr, + R2: tl.constexpr, + BLOCK_H: tl.constexpr, + WITH_INJECT_OUT: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + EARLY_TRIGGER: tl.constexpr, +): + pid = tl.program_id(0) + hh = pid * BLOCK_H + tl.arange(0, BLOCK_H) + sidx = tl.arange(0, HC) + col = sidx[None, :] * HIDDEN + hh[:, None] + hmask = (hh < HIDDEN)[:, None] & (sidx >= 0)[None, :] + # The whole [BLOCK_H, HC, R] weight tile is requested before any arithmetic. + ra = tl.arange(0, R1) + w_a = tl.load( + up_ptr + col[:, :, None] * stride_up_n + ra[None, None, :] * stride_up_r, + mask=hmask[:, :, None], + other=0.0, + ) + if R2 > 0: + rb = R1 + tl.arange(0, R2) + w_b = tl.load( + up_ptr + col[:, :, None] * stride_up_n + rb[None, None, :] * stride_up_r, + mask=hmask[:, :, None], + other=0.0, + ) + gam = 1.0 + tl.load(norm_w_ptr + col, mask=hmask, other=0.0).to(tl.float32) + if LAUNCH_WITH_PDL: + if EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() + # Weights above stream while the producer runs; its outputs are read after this. + tl.extra.cuda.gdc_wait() + for m in tl.static_range(M): + ssm = tl.load(ss_ptr + sidx * stride_ss + m) + pa = tl.load(part_ptr + sidx[:, None] * stride_ps + m * stride_pm + ra[None, :]) + if R2 > 0: + pb = tl.load(part_ptr + sidx[:, None] * stride_ps + m * stride_pm + rb[None, :]) + x = tl.load(x_ptr + m * stride_x + col, mask=hmask, other=0.0).to(tl.float32) + rstd = tl.rsqrt(ssm / HIDDEN + eps) + za = tl.sum(pa * rstd[:, None], axis=0) * (1.0 / HC) + acc = tl.sum(w_a.to(tl.float32) * (za * tl.sigmoid(za))[None, None, :], axis=2) + if R2 > 0: + zb = tl.sum(pb * rstd[:, None], axis=0) * (1.0 / HC) + acc += tl.sum(w_b.to(tl.float32) * (zb * tl.sigmoid(zb))[None, None, :], axis=2) + v = tl.sigmoid(acc) * (x * rstd[None, :] * gam) + tl.store( + mixed_ptr + m * stride_mixed + hh, + (tl.sum(v, axis=1) * (1.0 / HC)).to(mixed_ptr.dtype.element_ty), + mask=hh < HIDDEN, + ) + if WITH_INJECT_OUT: + ci = tl.arange(0, HC) + pi = tl.load( + part_ptr + sidx[:, None] * stride_ps + m * stride_pm + (R1 + R2) + ci[None, :], + mask=(sidx >= 0)[:, None] & (pid == 0), + other=0.0, + ) + zi = tl.sum(pi * rstd[:, None], axis=0) * (1.0 / HC) + tl.store( + inj_out_ptr + m * HC + ci, + (2.0 * tl.sigmoid(zi)).to(inj_out_ptr.dtype.element_ty), + mask=(ci >= 0) & (pid == 0), + ) + + +#: Launch configuration: (BLOCK_H, num_warps). Swept on GB300 (see the note). +DEFAULT_CONFIG = (8, 8) +#: Let the kernel after this one launch as soon as every program has requested its +#: weights: a programmatic dependent (the decode projections) then streams its own +#: weight while this kernel finishes. Its `gdc_wait` still waits for this whole grid. +TRIGGER_DEPENDENTS = True + + +def _rank_split(rank: int) -> Tuple[int, int]: + """`rank = R1 + R2` with `R1` the largest power of two not above it; `R2` must be + a power of two or zero (320 = 256 + 64).""" + r1 = 1 << (rank.bit_length() - 1) + return r1, rank - r1 + + +def launch( + part, + ss, + up, + hc_out, + norm_weight, + eps, + hc_count, + with_inject, + config=DEFAULT_CONFIG, + trigger=TRIGGER_DEPENDENTS, +): + """The launch behind `Qwen4ExpHcDecodeMix.__call__`, tile config exposed. + Returns `(mixed, inject_out)`.""" + rows, width = hc_out.shape + hidden = width // hc_count + rank = up.shape[0] + r1, r2 = _rank_split(rank) + block_h, warps = config + dev, dt = hc_out.device, hc_out.dtype + mixed = torch.empty((rows, hidden), dtype=dt, device=dev) + inject_out = torch.empty((rows, hc_count), dtype=dt, device=dev) if with_inject else None + pdl = pdl_enabled(dev.index if dev.index is not None else torch.cuda.current_device()) + if rows > 0: + _hc_decode_mix_kernel[(triton.cdiv(hidden, block_h),)]( + part, + ss, + up, + hc_out, + norm_weight, + mixed, + inject_out if with_inject else mixed, + part.stride(0), + part.stride(1), + ss.stride(0), + up.stride(0), + up.stride(1), + hc_out.stride(0), + mixed.stride(0), + eps, + M=rows, + HIDDEN=hidden, + HC=hc_count, + R1=r1, + R2=r2, + BLOCK_H=block_h, + WITH_INJECT_OUT=with_inject, + LAUNCH_WITH_PDL=pdl, + EARLY_TRIGGER=trigger, + num_warps=warps, + launch_pdl=pdl, + ) + return mixed, inject_out + + +class Qwen4ExpHcDecodeMix(OpWrapper): + """With `C = hc_count`, `H = width // C`, `R = up.shape[0]`, `part` / `ss` the + per-stream partial logits `[C, T, >= R (+ C)]` and sums of squares `[C, T]` + from `qwen4_exp_hc_decode_first`, `up` the `[R, width]` view of the + `[width, R]` up-projection weight:: + + rstd[t, s] = 1 / sqrt(ss[s, t] / H + eps) + logits[t, n] = sum_s rstd[t, s] * part[s, t, n] + low = silu(logits[:, :R] / C) + gate = sigmoid(low @ up) # [T, C*H] + normed[t, s*H+i] = hc_out[t, s*H+i] * rstd[t, s] * (1 + norm_weight[s*H+i]) + mixed[t, i] = mean over s of gate[t, s*H+i] * normed[t, s*H+i] + inject_out = 2 * sigmoid(logits[:, R:R+C] / C) # with_inject only + + fp32 inside, one rounding per output. Returns fresh `(mixed, inject_out)` + in `hc_out`'s dtype; `inject_out` is None without `with_inject`. + + Fusion boundary: the grouped RMS, the bottleneck SiLU, the injection + weights, the up projection, the sigmoid gate, the normed stream and the + stream mean. The partials are the producer's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: rows, with_inject (True), hc_count (4), hidden (2560), rank (320), dtype (bf16). + CELLS: tuple[Cell, ...] = tuple( + Cell(why=f"checkpoint 4x2560 r320 bf16, inject, {why}", spec=dict(rows=t)) + for t, why in ( + (1, "decode 1 row (con 1)"), + (2, "decode 2 rows"), + (3, "decode 3 rows"), + (8, "decode 8 rows"), + (16, "decode 16 rows (MAX_ROWS)"), + ) + ) + ( + Cell(why="stack-final mixer (no inject), 1 row", spec=dict(rows=1, with_inject=False)), + Cell(why="stack-final mixer (no inject), 16 rows", spec=dict(rows=16, with_inject=False)), + Cell(why="hc_out a row slice buf[::2], 4 rows", spec=dict(rows=4, layout="row_slice")), + Cell(why="hc_count 2 at hidden 2560, 2 rows", spec=dict(rows=2, hc_count=2)), + Cell( + why="hidden 128 rank 32 (not the checkpoint), 5 rows", + spec=dict(rows=5, hidden=128, rank=32), + ), + Cell(why="float16 operands, 4 rows", spec=dict(rows=4, dtype=torch.float16)), + Cell(why="zero rows", spec=dict(rows=0)), + ) + + note = """ + Metadata consumed: none. Call sites: `DecodeTarget`'s trunk, right after + each `qwen4_exp_hc_decode_first` (97 per forward, 1-row decode batches). + + Preconditions: `part` fp32 `[C, T, >= R (+ C)]` and `ss` fp32 `[C, T]`, + both with unit last stride (other strides free); `up` `[R, width]` (any + strides; the `.t()` view of the stored weight reads rows of `R` contiguous + elements); `R` a power of two or a sum of two (320 = 256 + 64); `hc_out` + `[T, width]` unit column stride; `norm_weight` contiguous `[width]`. Row + count is a compile-time constant, at most MAX_ROWS certified. + + Numerics: fp32 throughout (the composition it replaces rounded `low`, the + gate logits and `normed` to bf16), one rounding per output. Every program + recomputes `rstd` and `low` from the partials in the same order, so all + programs agree bitwise. Graph-capturable. Launched with programmatic + dependent launch on sm_90+ (`TRTLLM_ENABLE_PDL=0` turns it off): the + weight tile streams before `gdc_wait`, the producer's outputs after it. + + Grid: cdiv(H, 8) programs of 8 warps, each one [8, C, R] weight load + (20 KB). Measured on GB300 (graph replay, cold weights): ~4 us at 1 row; + each further row costs ~3 us (the row loop serializes its L2 round trips). + """ + + def __call__( + self, + part: torch.Tensor, + ss: torch.Tensor, + up: torch.Tensor, + hc_out: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + hc_count: int, + with_inject: bool, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + return launch(part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject) + + def reference( + self, + part: torch.Tensor, + ss: torch.Tensor, + up: torch.Tensor, + hc_out: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + hc_count: int, + with_inject: bool, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """float64, one rounding per output.""" + dtype = hc_out.dtype + rank = up.shape[0] + hidden = hc_out.shape[1] // hc_count + rstd = torch.rsqrt(ss.double().t() / hidden + eps) + logits = torch.einsum("ts,stn->tn", rstd, part.double()) + low = F.silu(logits[:, :rank] / hc_count) + gate = torch.sigmoid(low @ up.double()) + x = hc_out.double().unflatten(-1, (hc_count, hidden)) * rstd.unsqueeze(-1) + normed = x.flatten(-2) * (1.0 + norm_weight.double()) + mixed = (gate * normed).unflatten(-1, (hc_count, hidden)).mean(dim=-2).to(dtype) + inject_out = None + if with_inject: + inject_out = (2.0 * torch.sigmoid(logits[:, rank : rank + hc_count] / hc_count)).to( + dtype + ) + return mixed, inject_out + + def is_valid( + self, + part: torch.Tensor, + ss: torch.Tensor, + up: torch.Tensor, + hc_out: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + hc_count: int, + with_inject: bool, + ) -> None: + rows, width = hc_out.shape + rank = up.shape[0] + assert hc_out.stride(1) == 1, "hc_out unit column stride" + assert part.dtype == torch.float32 and part.stride(2) == 1, "part fp32, unit last stride" + assert part.shape[:2] == (hc_count, rows), "part [C, T, N]" + assert part.shape[2] >= rank + (hc_count if with_inject else 0), "part too narrow" + assert ss.dtype == torch.float32 and ss.shape == (hc_count, rows), "ss fp32 [C, T]" + assert rows <= 1 or ss.stride(1) == 1, "ss unit stride along rows" + assert up.shape == (rank, width), "up [R, width]" + assert norm_weight.shape == (width,) and norm_weight.is_contiguous() + r2 = _rank_split(rank)[1] + assert r2 & (r2 - 1) == 0, "rank must be a power of two or a sum of two" + + +qwen4_exp_hc_decode_mix = Qwen4ExpHcDecodeMix() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix_fp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix_fp4.py new file mode 100644 index 000000000000..4bb07863796a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix_fp4.py @@ -0,0 +1,431 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Decode-row hyper-connection site at an MoE site, second half: +`qwen4_exp_hc_decode_mix` (grouped RMS, SiLU bottleneck, injection weights, up +projection, sigmoid gate, stream mean) plus the NVFP4 quantization of the mean +for the routed experts -- one Triton GEMV launch. + +Tier 5 (Triton written for this target). Same structure as +`qwen4_exp_hc_decode_mix`, except that each program owns a whole number of +16-column NVFP4 blocks so it can quantize its own columns: the scale block is +the only cross-column dependency the quantization adds. + +The quantization helpers here (`nvfp4_quantize_blocks` in Triton, +`nvfp4_reference` in torch) reproduce `torch.ops.trtllm.fp4_quantize` (NVFP4, +linear scales) step for step, including its two `rcp.approx.ftz.f32` +reciprocals, so the packed data and the scales are bitwise that op's output on +the rounded mix; `qwen4_exp_hc_gate_mean_fp4` imports them. +""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .qwen4_exp_hc_decode_mix import _rank_split, pdl_enabled + +#: Largest row count this entry is certified for. +MAX_ROWS = 16 + +#: NVFP4 block: one e4m3 scale per this many contiguous elements. +SF_VEC = 16 +_E4M3_MAX = 448.0 +_E2M1_MAX = 6.0 +_E2M1_MIDPOINTS = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0) +_TIE_UP = (False, True, False, True, False, True, False) +#: Same window `catalog/quantization/fp4_quantize` states for the same two +#: approximate reciprocals. +TIE_WINDOW = 2.0**-20 + + +@triton.jit +def rcp_approx_ftz(x): + """`rcp.approx.ftz.f32`, the reciprocal the trtllm quantizer uses.""" + return tl.inline_asm_elementwise( + "rcp.approx.ftz.f32 $0, $1;", "=r,r", [x], dtype=tl.float32, is_pure=True, pack=1 + ) + + +@triton.jit +def nvfp4_quantize_blocks(v, gs): + """NVFP4 of `v` `[NB, 8, 2]` fp32 (values already rounded to the activation + dtype; element `16*b + 2*p + j` of the row chunk at `[b, p, j]`) with the + global scale `gs` `[NB]`. Returns the packed bytes `[NB, 8]` (int32, element + `2k` in the low nibble) and the e4m3 scale bytes `[NB]` (uint8).""" + amax = tl.max(tl.max(tl.abs(v), axis=2), axis=1) + sfv = gs * (amax * rcp_approx_ftz(tl.zeros_like(amax) + 6.0)) + sf8 = sfv.to(tl.float8e4nv) + osc = rcp_approx_ftz(sf8.to(tl.float32) * rcp_approx_ftz(gs)) + osc = tl.where(amax != 0.0, osc, 0.0) + s = v * osc[:, None, None] + a = tl.abs(s) + # Round to nearest, ties to the even e2m1 code; saturate at 6. + code = ( + (a > 0.25).to(tl.int32) + + (a >= 0.75).to(tl.int32) + + (a > 1.25).to(tl.int32) + + (a >= 1.75).to(tl.int32) + + (a > 2.5).to(tl.int32) + + (a >= 3.5).to(tl.int32) + + (a > 5.0).to(tl.int32) + ) + code = code | tl.where(s.to(tl.int32, bitcast=True) < 0, 8, 0) + shift = 4 * tl.arange(0, 2) + packed = tl.sum(code << shift[None, None, :], axis=2) + return packed, sf8.to(tl.uint8, bitcast=True) + + +def nvfp4_reference( + x: torch.Tensor, global_scale: torch.Tensor +) -> Tuple[torch.Tensor, torch.Tensor]: + """Native-torch NVFP4 of `x` `[T, K]` with linear scales: per 16-block, + `sf = e4m3(gs * amax / 6)` and the data `x * gs / sf` rounded to e2m1 + (nearest, ties to even). Returns `(packed [T, K/2] uint8, sf [T*K/16] uint8)`.""" + packed, sf, _ = _nvfp4_torch(x, global_scale) + return packed, sf + + +def nvfp4_near_tie(x: torch.Tensor, global_scale: torch.Tensor) -> torch.Tensor: + """`[T, K]` bool: elements whose pre-rounding value sits within `TIE_WINDOW` + of an e2m1 midpoint, where approximate reciprocals may pick either side.""" + scaled = _nvfp4_torch(x, global_scale)[2] + mid = torch.tensor(_E2M1_MIDPOINTS, device=x.device, dtype=torch.float32) + a = scaled.abs().unsqueeze(-1) + return ((a - mid).abs() <= TIE_WINDOW * a).any(-1) + + +def _nvfp4_torch(x, global_scale): + gs = global_scale.float().reshape(()) + rows, k = x.shape + blocks = x.float().reshape(rows, k // SF_VEC, SF_VEC) + vecmax = blocks.abs().amax(-1) + sf = (gs * (vecmax / _E2M1_MAX)).clamp(max=_E4M3_MAX).to(torch.float8_e4m3fn) + out_scale = torch.where(vecmax != 0, gs / sf.float(), torch.zeros_like(vecmax)) + scaled = (blocks * out_scale.unsqueeze(-1)).reshape(rows, k) + mid = torch.tensor(_E2M1_MIDPOINTS, device=x.device, dtype=torch.float32) + tie_up = torch.tensor(_TIE_UP, device=x.device, dtype=torch.bool) + a = scaled.abs().unsqueeze(-1) + code = ((a > mid).sum(-1) + ((a == mid) & tie_up).any(-1)).to(torch.uint8) + negative = (scaled < 0) | ((scaled == 0) & torch.signbit(scaled)) + code = code | (negative.to(torch.uint8) << 3) + packed = code[:, 0::2] | (code[:, 1::2] << 4) + return packed, sf.view(torch.uint8).reshape(-1), scaled + + +@triton.jit +def _hc_decode_mix_fp4_kernel( + part_ptr, + ss_ptr, + up_ptr, + x_ptr, + norm_w_ptr, + gs_ptr, + mixed_ptr, + inj_out_ptr, + q_ptr, + sf_ptr, + stride_ps, + stride_pm, + stride_ss, + stride_up_r, + stride_up_n, + stride_x, + stride_mixed, + eps, + M: tl.constexpr, + HIDDEN: tl.constexpr, + HC: tl.constexpr, + R1: tl.constexpr, + R2: tl.constexpr, + BLOCK_H: tl.constexpr, + WITH_INJECT_OUT: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + EARLY_TRIGGER: tl.constexpr, +): + NBH: tl.constexpr = BLOCK_H // 16 + pid = tl.program_id(0) + hh = pid * BLOCK_H + tl.arange(0, BLOCK_H) + sidx = tl.arange(0, HC) + col = sidx[None, :] * HIDDEN + hh[:, None] + hmask = (hh < HIDDEN)[:, None] & (sidx >= 0)[None, :] + ra = tl.arange(0, R1) + w_a = tl.load( + up_ptr + col[:, :, None] * stride_up_n + ra[None, None, :] * stride_up_r, + mask=hmask[:, :, None], + other=0.0, + ) + if R2 > 0: + rb = R1 + tl.arange(0, R2) + w_b = tl.load( + up_ptr + col[:, :, None] * stride_up_n + rb[None, None, :] * stride_up_r, + mask=hmask[:, :, None], + other=0.0, + ) + gam = 1.0 + tl.load(norm_w_ptr + col, mask=hmask, other=0.0).to(tl.float32) + b = pid * NBH + tl.arange(0, NBH) + p = tl.arange(0, 8) + bmask = b < HIDDEN // 16 + gs = tl.load(gs_ptr + tl.zeros([NBH], dtype=tl.int32)) + if LAUNCH_WITH_PDL: + if EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() + tl.extra.cuda.gdc_wait() + for m in tl.static_range(M): + ssm = tl.load(ss_ptr + sidx * stride_ss + m) + pa = tl.load(part_ptr + sidx[:, None] * stride_ps + m * stride_pm + ra[None, :]) + if R2 > 0: + pb = tl.load(part_ptr + sidx[:, None] * stride_ps + m * stride_pm + rb[None, :]) + x = tl.load(x_ptr + m * stride_x + col, mask=hmask, other=0.0).to(tl.float32) + rstd = tl.rsqrt(ssm / HIDDEN + eps) + za = tl.sum(pa * rstd[:, None], axis=0) * (1.0 / HC) + acc = tl.sum(w_a.to(tl.float32) * (za * tl.sigmoid(za))[None, None, :], axis=2) + if R2 > 0: + zb = tl.sum(pb * rstd[:, None], axis=0) * (1.0 / HC) + acc += tl.sum(w_b.to(tl.float32) * (zb * tl.sigmoid(zb))[None, None, :], axis=2) + v = tl.sigmoid(acc) * (x * rstd[None, :] * gam) + mixed = (tl.sum(v, axis=1) * (1.0 / HC)).to(mixed_ptr.dtype.element_ty) + tl.store(mixed_ptr + m * stride_mixed + hh, mixed, mask=hh < HIDDEN) + packed, sf = nvfp4_quantize_blocks(tl.reshape(mixed.to(tl.float32), [NBH, 8, 2]), gs) + q_off = m * (HIDDEN // 2) + b[:, None] * 8 + p[None, :] + tl.store(q_ptr + q_off, packed.to(tl.uint8), mask=bmask[:, None]) + tl.store(sf_ptr + m * (HIDDEN // 16) + b, sf, mask=bmask) + if WITH_INJECT_OUT: + ci = tl.arange(0, HC) + pi = tl.load( + part_ptr + sidx[:, None] * stride_ps + m * stride_pm + (R1 + R2) + ci[None, :], + mask=(sidx >= 0)[:, None] & (pid == 0), + other=0.0, + ) + zi = tl.sum(pi * rstd[:, None], axis=0) * (1.0 / HC) + tl.store( + inj_out_ptr + m * HC + ci, + (2.0 * tl.sigmoid(zi)).to(inj_out_ptr.dtype.element_ty), + mask=(ci >= 0) & (pid == 0), + ) + + +#: Launch configuration: (BLOCK_H, num_warps). BLOCK_H a multiple of 16. Swept +#: on GB300 at 1/2/4 rows (graph, cold weights): (16, 4) fastest at every row +#: count; (16, 2), (16, 8) and (32, 4) 0.3..5 us slower. +DEFAULT_CONFIG = (16, 4) +#: Let the kernel after this one launch as soon as every program has requested its +#: weights: a programmatic dependent (the decode projections) then streams its own +#: weight while this kernel finishes. Its `gdc_wait` still waits for this whole grid. +TRIGGER_DEPENDENTS = True + + +def launch( + part, + ss, + up, + hc_out, + norm_weight, + eps, + hc_count, + with_inject, + global_scale, + config=DEFAULT_CONFIG, + trigger=TRIGGER_DEPENDENTS, +): + """The launch behind `Qwen4ExpHcDecodeMixFp4.__call__`, tile config exposed. + Returns `(mixed, inject_out, xq, xsf)`.""" + rows, width = hc_out.shape + hidden = width // hc_count + rank = up.shape[0] + r1, r2 = _rank_split(rank) + block_h, warps = config + dev, dt = hc_out.device, hc_out.dtype + mixed = torch.empty((rows, hidden), dtype=dt, device=dev) + inject_out = torch.empty((rows, hc_count), dtype=dt, device=dev) if with_inject else None + xq = torch.empty((rows, hidden // 2), dtype=torch.uint8, device=dev) + xsf = torch.empty((rows * (hidden // SF_VEC),), dtype=torch.uint8, device=dev) + pdl = pdl_enabled(dev.index if dev.index is not None else torch.cuda.current_device()) + if rows > 0: + _hc_decode_mix_fp4_kernel[(triton.cdiv(hidden, block_h),)]( + part, + ss, + up, + hc_out, + norm_weight, + global_scale, + mixed, + inject_out if with_inject else mixed, + xq, + xsf, + part.stride(0), + part.stride(1), + ss.stride(0), + up.stride(0), + up.stride(1), + hc_out.stride(0), + mixed.stride(0), + eps, + M=rows, + HIDDEN=hidden, + HC=hc_count, + R1=r1, + R2=r2, + BLOCK_H=block_h, + WITH_INJECT_OUT=with_inject, + LAUNCH_WITH_PDL=pdl, + EARLY_TRIGGER=trigger, + num_warps=warps, + launch_pdl=pdl, + ) + return mixed, inject_out, xq, xsf + + +class Qwen4ExpHcDecodeMixFp4(OpWrapper): + """`qwen4_exp_hc_decode_mix` plus the NVFP4 of its `mixed`. With `C = + hc_count`, `H = width // C`, `R = up.shape[0]`, `part` / `ss` the per-stream + partial logits `[C, T, >= R (+ C)]` and sums of squares `[C, T]` from + `qwen4_exp_hc_decode_first`, `up` the `[R, width]` view of the up weight:: + + rstd[t, s] = 1 / sqrt(ss[s, t] / H + eps) + logits[t, n] = sum_s rstd[t, s] * part[s, t, n] + low = silu(logits[:, :R] / C) + gate = sigmoid(low @ up) # [T, C*H] + normed[t, s*H+i] = hc_out[t, s*H+i] * rstd[t, s] * (1 + norm_weight[s*H+i]) + mixed[t, i] = mean over s of gate[t, s*H+i] * normed[t, s*H+i] + inject_out = 2 * sigmoid(logits[:, R:R+C] / C) # with_inject only + xq, xsf = NVFP4(mixed) # e2m1 pairs, e4m3 scale per 16, linear layout + + fp32 inside, one rounding per output; the quantization reads the rounded + `mixed`. Returns fresh `(mixed, inject_out, xq [T, H/2] uint8, xsf [T*H/16] + uint8)`; `inject_out` is None without `with_inject`. `global_scale` is the + 1-element fp32 NVFP4 global scale. + + Fusion boundary: everything `qwen4_exp_hc_decode_mix` fuses, plus the NVFP4 + quantization of `mixed`. The partials are the producer's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: rows, with_inject (True), hc_count (4), hidden (2560), rank (320), dtype (bf16). + CELLS: tuple[Cell, ...] = tuple( + Cell(why=f"checkpoint 4x2560 r320 bf16, inject, {why}", spec=dict(rows=t)) + for t, why in ( + (1, "decode 1 row (con 1)"), + (2, "decode 2 rows"), + (3, "decode 3 rows"), + (16, "decode 16 rows (MAX_ROWS)"), + ) + ) + ( + Cell(why="hc_out a row slice buf[::2], 4 rows", spec=dict(rows=4, layout="row_slice")), + Cell(why="hc_count 2 at hidden 2560, 2 rows", spec=dict(rows=2, hc_count=2)), + Cell( + why="hidden 144 rank 32 (9 scale blocks), 5 rows", + spec=dict(rows=5, hidden=144, rank=32), + ), + Cell(why="no inject, 1 row", spec=dict(rows=1, with_inject=False)), + Cell(why="float16 operands, 4 rows", spec=dict(rows=4, dtype=torch.float16)), + Cell(why="zero rows", spec=dict(rows=0)), + ) + + note = """ + Metadata consumed: none. Call sites: `DecodeTarget`'s trunk, the MoE-site + mixer of every layer (48 per forward) at batches of at most + `_HC_DECODE_GEMV_MAX_ROWS` rows, right after `qwen4_exp_hc_decode_first`; + `xq`/`xsf` feed `fp4_block_scale_moe_renorm_routed` (linear scales). + + Preconditions: those of `qwen4_exp_hc_decode_mix`, plus `H` a multiple of + 16 and `global_scale` one fp32 element on the device. Row count is a + compile-time constant, at most MAX_ROWS certified. + + Numerics: `mixed` and `inject_out` follow `qwen4_exp_hc_decode_mix`'s + formula (fp32, one rounding); `inject_out` is bitwise its result, `mixed` + can differ in a few elements by one rounding because the 16-column program + tile reduces the up projection in another order. The NVFP4 is + the trtllm recipe (two `rcp.approx.ftz.f32`, `cvt.rn.satfinite` e4m3 + scale, e2m1 round-to-nearest-even saturating at 6) on the rounded `mixed`; + the GPU test checks it bitwise against `torch.ops.trtllm.fp4_quantize`. + Graph-capturable; programmatic dependent launch as `qwen4_exp_hc_decode_mix`. + + Grid: cdiv(H, BLOCK_H) programs (BLOCK_H a multiple of 16; 16 -> 160 + programs of 4 warps, each one [16, C, R] weight load, 40 KB). Measured on + GB300 (graph replay, cold weights, 1 row): 4.6 us, against 6.0 us for + `qwen4_exp_hc_decode_mix` + `fp4_quantize` and 3.9 us for + `qwen4_exp_hc_decode_mix` alone (its 8-column tile has twice the programs). + """ + + def __call__( + self, + part: torch.Tensor, + ss: torch.Tensor, + up: torch.Tensor, + hc_out: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + hc_count: int, + with_inject: bool, + global_scale: torch.Tensor, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, torch.Tensor]: + return launch(part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject, global_scale) + + def reference( + self, + part: torch.Tensor, + ss: torch.Tensor, + up: torch.Tensor, + hc_out: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + hc_count: int, + with_inject: bool, + global_scale: torch.Tensor, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, torch.Tensor]: + """float64 with one rounding per output, then the NVFP4 of the rounded + `mixed` in native torch.""" + dtype = hc_out.dtype + rank = up.shape[0] + hidden = hc_out.shape[1] // hc_count + rstd = torch.rsqrt(ss.double().t() / hidden + eps) + logits = torch.einsum("ts,stn->tn", rstd, part.double()) + low = F.silu(logits[:, :rank] / hc_count) + gate = torch.sigmoid(low @ up.double()) + x = hc_out.double().unflatten(-1, (hc_count, hidden)) * rstd.unsqueeze(-1) + normed = x.flatten(-2) * (1.0 + norm_weight.double()) + mixed = (gate * normed).unflatten(-1, (hc_count, hidden)).mean(dim=-2).to(dtype) + inject_out = None + if with_inject: + inject_out = (2.0 * torch.sigmoid(logits[:, rank : rank + hc_count] / hc_count)).to( + dtype + ) + return (mixed, inject_out) + nvfp4_reference(mixed, global_scale) + + def is_valid( + self, + part: torch.Tensor, + ss: torch.Tensor, + up: torch.Tensor, + hc_out: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + hc_count: int, + with_inject: bool, + global_scale: torch.Tensor, + ) -> None: + rows, width = hc_out.shape + rank = up.shape[0] + assert hc_out.stride(1) == 1, "hc_out unit column stride" + assert part.dtype == torch.float32 and part.stride(2) == 1, "part fp32, unit last stride" + assert part.shape[:2] == (hc_count, rows), "part [C, T, N]" + assert part.shape[2] >= rank + (hc_count if with_inject else 0), "part too narrow" + assert ss.dtype == torch.float32 and ss.shape == (hc_count, rows), "ss fp32 [C, T]" + assert rows <= 1 or ss.stride(1) == 1, "ss unit stride along rows" + assert up.shape == (rank, width), "up [R, width]" + assert norm_weight.shape == (width,) and norm_weight.is_contiguous() + r2 = _rank_split(rank)[1] + assert r2 & (r2 - 1) == 0, "rank must be a power of two or a sum of two" + assert (width // hc_count) % SF_VEC == 0, "hidden must be whole 16-blocks" + # The kernel loads element 0 only; a per-row scale would be applied as row 0's. + assert global_scale.numel() == 1 and global_scale.dtype == torch.float32 + + +qwen4_exp_hc_decode_mix_fp4 = Qwen4ExpHcDecodeMixFp4() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix_fp4_cute.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix_fp4_cute.py new file mode 100644 index 000000000000..a12ca131fe3b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix_fp4_cute.py @@ -0,0 +1,443 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Decode-row hyper-connection site at an MoE site, second half, for one-row batches: +the grouped RMS, the SiLU bottleneck and the injection weights from the first half's +partials, the up projection, the sigmoid gate, the stream mean and the NVFP4 of the +mean for the routed experts -- one PDL launch. + +Tier 5 (CuTe DSL written for this target). `qwen4_exp_hc_decode_mix_fp4` computes the +same for up to 16 rows in Triton, whose programs hold a 64-row x 320 tile of the up +weight in 128 threads. Here a CTA +owns one 16-column NVFP4 block of the hidden width across the 4 streams (64 rows of +the up weight), and 8 threads share a row, 40 weights each; 160 CTAs of 512 threads, +all resident at once. + +Order of work inside a CTA: + +1. its 64 weight rows and its norm-weight elements into registers (static), before + `griddepcontrol.wait`; +2. after the wait: the four sums of squares (every thread), one partial column per + thread for the first 324 threads -> the shared `low` vector and, in CTA 0, the + injection weights; each thread's residual element; +3. each row's dot with `low` over its 8 threads, the sigmoid gate on the normed + element, the mean over streams (16 threads), the store; with NVFP4, warp 0 also + quantizes the block as `torch.ops.trtllm.fp4_quantize` does (linear scales, + `rcp.approx.ftz` reciprocals). +""" + +from __future__ import annotations + +import functools + +import cuda.bindings.driver as cuda_drv +import cutlass +import cutlass.cute as cute +import torch +from cutlass import const_expr +from cutlass.cute import experimental as cute_ext +from cutlass.cute.runtime import from_dlpack + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .qwen4_exp_hc_decode_first_cute import pdl_enabled +from .qwen4_exp_hc_decode_mix import qwen4_exp_hc_decode_mix as _triton_mix +from .qwen4_exp_hc_decode_mix_fp4 import qwen4_exp_hc_decode_mix_fp4 as _triton_mix_fp4 + +MAX_ROWS = 1 +VEC = 8 +#: Hidden columns per CTA: one NVFP4 scale block. +BLOCK_H = 16 +#: Threads sharing one up-weight row. +GROUP = 8 +_COMPILE_OPTIONS = "--ptxas-options -maxrregcount=64" + + +class _MixKernel: + def __init__(self, *, hidden, streams, rank, inject_out, fp4, pdl): + assert rank % (GROUP * VEC) == 0 and hidden % BLOCK_H == 0 + self.hidden, self.streams, self.rank = hidden, streams, rank + self.inject_out, self.fp4, self.pdl = inject_out, fp4, pdl + self.k_tiles = rank // (GROUP * VEC) + self.threads = streams * BLOCK_H * GROUP + assert self.threads >= rank + streams and self.threads % 32 == 0 + + @cute.jit + def __call__( + self, + mPart, + mSs, + mUp, + mHc, + mNorm, + mGs, + mMixed, + mInj, + mXq, + mXsf, + eps: cutlass.Float32, + stream: cuda_drv.CUstream, + ): + load_static = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + cutlass.BFloat16, + num_bits_per_copy=16 * VEC, + load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING, + ) + self.kernel( + mPart, mSs, mUp, mHc, mNorm, mGs, mMixed, mInj, mXq, mXsf, eps, load_static + ).launch( + grid=[self.hidden // BLOCK_H, 1, 1], + block=[self.threads, 1, 1], + smem=((self.rank // VEC) * (VEC + 1) + self.streams * BLOCK_H) * 4 + 32, + stream=stream, + use_pdl=self.pdl, + min_blocks_per_mp=1, + ) + + @cute.kernel + def kernel( + self, + mPart, + mSs, + mUp, + mHc, + mNorm, + mGs, + mMixed, + mInj, + mXq, + mXsf, + eps: cutlass.Float32, + load_static: cute.CopyAtom, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + warp = cute.arch.warp_idx() + lane = cute.arch.lane_idx() + hidden: cutlass.Constexpr = self.hidden + streams: cutlass.Constexpr = self.streams + rank: cutlass.Constexpr = self.rank + k_tiles: cutlass.Constexpr = self.k_tiles + i0 = bidx * BLOCK_H + row = tidx // GROUP # (stream, column) row of this CTA: stream-major + q = tidx % GROUP + s = row // BLOCK_H + i = row % BLOCK_H + j = s * hidden + i0 + i # the up-weight row, a column of the residual width + + # 1. Static operands, before the wait. + u_vec = cute.logical_divide(mUp, (None, VEC)) + u_reg = cute.make_rmem_tensor( + cute.make_layout((k_tiles, VEC), stride=(VEC, 1)), cutlass.BFloat16 + ) + for kt in cutlass.range_constexpr(k_tiles): + cute.copy(load_static, u_vec[j, (None, kt * GROUP + q)], u_reg[kt, None]) + gam = 1.0 + mNorm[0, j].to(cutlass.Float32) + gs = cutlass.Float32(0.0) + if const_expr(self.fp4): + gs = mGs[0] + if const_expr(self.pdl): + cute.arch.griddepcontrol_launch_dependents() + cute.arch.griddepcontrol_wait() + + # 2. Grouped RMS, logits, low (shared), injection weights. Every producer value this + # thread needs is requested before the first is used (a partial column clamped into + # range for the threads past it), so the CTA pays one L2 round trip. + cols: cutlass.Constexpr = rank + (streams if self.inject_out else 0) # partial columns read + pcol = cutlass.min(tidx, cols - 1) + pv = [mPart[t, pcol] for t in range(streams)] + ssv = [mSs[t] for t in range(streams)] + x = mHc[0, j].to(cutlass.Float32) + rstd = cute.make_rmem_tensor(cute.make_layout((streams,), stride=(1,)), cutlass.Float32) + for t in cutlass.range_constexpr(streams): + rstd[t] = cute.math.rsqrt(ssv[t] / hidden + eps) + smem = cutlass.utils.SmemAllocator() + # `low` in rows of VEC with one pad word, so the 8 threads of a row (stride VEC + 1 + # words) read 8 different banks. + low = smem.allocate_tensor( + cutlass.Float32, + cute.make_layout((rank // VEC, VEC), stride=(VEC + 1, 1)), + byte_alignment=16, + ) + vals = smem.allocate_tensor( + cutlass.Float32, + cute.make_layout((streams, BLOCK_H), stride=(BLOCK_H, 1)), + byte_alignment=16, + ) + if tidx < cols: + logit = cutlass.Float32(0.0) + for t in cutlass.range_constexpr(streams): + logit = pv[t] * rstd[t] + logit + z = logit * (1.0 / streams) + sig = 1.0 / (1.0 + cute.math.exp(-z)) + if tidx < rank: + low[tidx // VEC, tidx % VEC] = z * sig + else: + if const_expr(self.inject_out): + if bidx == 0: + mInj[0, tidx - rank] = (2.0 * sig).to(cutlass.BFloat16) + cute.arch.sync_threads() + + # 3. Up projection over the row's 8 threads (two independent sums), gate, normed + # element, stream mean. + sums = [cutlass.Float32(0.0), cutlass.Float32(0.0)] + for kt in cutlass.range_constexpr(k_tiles): + for e in cutlass.range_constexpr(VEC): + sums[e % 2] = ( + u_reg[kt, e].to(cutlass.Float32) * low[kt * GROUP + q, e] + sums[e % 2] + ) + dot = sums[0] + sums[1] + dot = cute.arch.warp_reduction(dot, lambda a, b: a + b, threads_in_group=GROUP) + if q == 0: + gate = 1.0 / (1.0 + cute.math.exp(-dot)) + vals[s, i] = gate * (x * rstd[s] * gam) + cute.arch.sync_threads() + if warp == 0: + m = cutlass.Float32(0.0) + if lane < BLOCK_H: + tot = cutlass.Float32(0.0) + for t in cutlass.range_constexpr(streams): + tot = tot + vals[t, lane] + mb = (tot * (1.0 / streams)).to(cutlass.BFloat16) + mMixed[0, i0 + lane] = mb + m = mb.to(cutlass.Float32) + if const_expr(self.fp4): + # NVFP4 of the 16 rounded values in lanes 0..15 (lanes 16..31 a second, + # discarded group), as `fp4_quantize`: e4m3(gs * amax / 6), then + # x * rcp(sf * rcp(gs)) rounded to e2m1, nearest, ties to even. + amax = cute.arch.warp_reduction_max(cute.math.abs(m), threads_in_group=BLOCK_H) + sf8 = (gs * (amax * cute.arch.rcp_approx(6.0))).to(cutlass.Float8E4M3FN) + osc = cute.arch.rcp_approx(sf8.to(cutlass.Float32) * cute.arch.rcp_approx(gs)) + if amax == 0.0: + osc = cutlass.Float32(0.0) + sv = m * osc + a = cute.math.abs(sv) + code = cutlass.Int32(0) + if a > 0.25: + code = code + 1 + if a >= 0.75: + code = code + 1 + if a > 1.25: + code = code + 1 + if a >= 1.75: + code = code + 1 + if a > 2.5: + code = code + 1 + if a >= 3.5: + code = code + 1 + if a > 5.0: + code = code + 1 + if cute.math.copysign(1.0, sv) < 0.0: # the sign bit, -0.0 included + code = code + 8 + hi = cute.arch.shuffle_sync_down(code, offset=1) + if lane < BLOCK_H: + if lane % 2 == 0: + mXq[0, (i0 + lane) // 2] = (code + hi * 16).to(cutlass.Uint8) + if lane == 0: + mXsf[bidx] = sf8 + + +def _dl(t: torch.Tensor, align: int = 16): + return from_dlpack(t, assumed_align=align) + + +@functools.lru_cache(maxsize=None) +def _compiled( + hidden, streams, rank, part_cols, part_stride, ss_stride, hc_stride, inject_out, fp4, pdl +): + kern = _MixKernel( + hidden=hidden, streams=streams, rank=rank, inject_out=inject_out, fp4=fp4, pdl=pdl + ) + dev, bf = "cuda", torch.bfloat16 + width = hidden * streams + part = torch.empty(streams, part_stride, dtype=torch.float32, device=dev)[:, :part_cols] + ss = torch.empty(streams * ss_stride, dtype=torch.float32, device=dev)[::ss_stride] + up = torch.empty(width, rank, dtype=bf, device=dev) + hc = torch.empty(1, hc_stride, dtype=bf, device=dev)[:, :width] + norm = torch.empty(1, width, dtype=bf, device=dev) + gs = torch.empty(1, dtype=torch.float32, device=dev) + mixed = torch.empty(1, hidden, dtype=bf, device=dev) + inj = torch.empty(1, streams, dtype=bf, device=dev) + xq = torch.empty(1, hidden // 2, dtype=torch.uint8, device=dev) + xsf = torch.empty(hidden // 16, dtype=torch.float8_e4m3fn, device=dev) + stream = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream) + return cute_ext.compile( + kern, + _dl(part, 4), + _dl(ss, 4), + _dl(up), + _dl(hc, 2), + _dl(norm, 2), + _dl(gs, 4), + _dl(mixed, 2), + _dl(inj, 2), + _dl(xq, 1), + _dl(xsf, 1), + cutlass.Float32(1e-6), + stream, + options=_COMPILE_OPTIONS, + ) + + +def launch(part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject, global_scale=None): + """The launch behind `Qwen4ExpHcDecodeMixFp4Cute.__call__`. Returns `(mixed, + inject_out, xq, xsf)`; without `global_scale` (the kernel's plain build, kept for + measurement) `xq` / `xsf` are None.""" + rows, width = hc_out.shape + hidden = width // hc_count + rank = up.shape[0] + dev, dt = hc_out.device, hc_out.dtype + fp4 = global_scale is not None + mixed = torch.empty((rows, hidden), dtype=dt, device=dev) + inject_out = torch.empty((rows, hc_count), dtype=dt, device=dev) if with_inject else None + xq = torch.empty((rows, hidden // 2), dtype=torch.uint8, device=dev) if fp4 else None + xsf = torch.empty((rows * (hidden // 16),), dtype=torch.uint8, device=dev) if fp4 else None + if rows == 0: + return mixed, inject_out, xq, xsf + pdl = pdl_enabled(dev.index if dev.index is not None else torch.cuda.current_device()) + part2 = part[:, 0, :] + fn = _compiled( + hidden, + hc_count, + rank, + part2.shape[1], + part2.stride(0), + ss.stride(0), + hc_out.stride(0), + bool(with_inject), + fp4, + bool(pdl), + ) + stream = cuda_drv.CUstream(torch.cuda.current_stream(dev).cuda_stream) + gs = global_scale.reshape(1) if fp4 else mixed.new_empty(1, dtype=torch.float32) + inj = inject_out if with_inject else mixed.new_empty(1, hc_count) + q = xq if fp4 else mixed.new_empty(1, hidden // 2, dtype=torch.uint8) + sf = ( + xsf.view(torch.float8_e4m3fn) + if fp4 + else mixed.new_empty(hidden // 16, dtype=torch.float8_e4m3fn) + ) + fn( + _dl(part2, 4), + _dl(ss[:, 0], 4), + _dl(up.t()), + _dl(hc_out, 2), + _dl(norm_weight.reshape(1, -1), 2), + _dl(gs, 4), + _dl(mixed, 2), + _dl(inj, 2), + _dl(q, 1), + _dl(sf, 1), + cutlass.Float32(eps), + stream, + ) + return mixed, inject_out, xq, xsf + + +def _check(part, ss, up, hc_out, norm_weight, hc_count) -> None: + rows, width = hc_out.shape + assert rows <= MAX_ROWS, "one row" + assert width // hc_count == 2560 and up.shape[0] == 320 and hc_count == 4, ( + "the checkpoint's 4 x 2560, rank 320" + ) + assert hc_out.dtype == torch.bfloat16 and up.dtype == torch.bfloat16, "bf16" + assert up.t().is_contiguous(), "up the [R, width] view of a row-major [width, R] weight" + + +_CELLS = ( + Cell(why="checkpoint 4x2560 r320 bf16, inject, 1 row (con 1)", spec=dict(rows=1)), + Cell(why="stack-final mixer (no inject), 1 row", spec=dict(rows=1, with_inject=False)), + Cell(why="hc_out a row slice buf[::2], 1 row", spec=dict(rows=1, layout="row_slice")), +) + + +class Qwen4ExpHcDecodeMixFp4Cute(OpWrapper): + """`qwen4_exp_hc_decode_mix_fp4` for a one-row batch: `(mixed, inject_out, xq, + xsf)`, the NVFP4 of the rounded `mixed` with linear scales (see that entry). + + Fusion boundary: as `qwen4_exp_hc_decode_mix_fp4`. + """ + + ARCHS = frozenset({Arch.SM_103}) + CELLS: tuple[Cell, ...] = _CELLS + + note = """ + Metadata consumed: none. Call sites: `DecodeTarget`'s MoE sites at a one-row + decode step, after `qwen4_exp_hc_decode_first_cute`; batches of more rows keep the + Triton entries. + + Preconditions: the Triton entry's, and one row, the checkpoint's 4 x 2560 width, + rank 320, bf16; `up` the `.t()` view of a contiguous `[width, R]` weight; + `global_scale` fp32 with one element. + + Numerics: fp32 inside, one rounding per output, the Triton entry's formula and + order up to fp32 reassociation of the up-projection dot (inside its band); the + NVFP4 is `fp4_quantize`'s arithmetic on the rounded `mixed` (`rcp.approx.ftz` + reciprocals, e4m3 scale rounded to nearest, e2m1 nearest-even with saturation at + 6). Deterministic; graph-capturable once compiled (the first call per variant). + PDL: the weight rows and norm elements stream before `griddepcontrol.wait`. + """ + + def __call__(self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject, global_scale): + return launch(part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject, global_scale) + + def reference(self, *args, **kwargs): + """The Triton entry's reference: float64, then the NVFP4 of the rounded `mixed`.""" + return _triton_mix_fp4.reference(*args, **kwargs) + + def is_valid( + self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject, global_scale + ) -> None: + _triton_mix_fp4.is_valid( + part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject, global_scale + ) + _check(part, ss, up, hc_out, norm_weight, hc_count) + assert global_scale.numel() == 1 and global_scale.dtype == torch.float32, ( + "global_scale fp32 [1]" + ) + + +qwen4_exp_hc_decode_mix_fp4_cute = Qwen4ExpHcDecodeMixFp4Cute() + + +class Qwen4ExpHcDecodeMixCute(OpWrapper): + """`qwen4_exp_hc_decode_mix` for a one-row batch: `(mixed, inject_out)` -- this + kernel's plain build (no NVFP4 output). + + Fusion boundary: as `qwen4_exp_hc_decode_mix`. + """ + + ARCHS = frozenset({Arch.SM_103}) + CELLS: tuple[Cell, ...] = _CELLS + + note = """ + Metadata consumed: none. Call sites: `DecodeTarget`'s non-MoE sites at a one-row + decode step, after `qwen4_exp_hc_decode_first_cute`; batches of more rows keep the + Triton entry. Beside a first kernel capped at 96 registers, its 160 CTAs (512 + threads at 49 registers) are all resident while that kernel runs, so the up weight + is in registers when the wait returns; the Triton entry's 320 programs are not. + + Preconditions: the Triton entry's, and one row, the checkpoint's 4 x 2560 width, + rank 320, bf16; `up` the `.t()` view of a contiguous `[width, R]` weight. + + Numerics: fp32 inside, one rounding per output, the Triton entry's formula and + order up to fp32 reassociation of the up-projection dot (inside its band). + Deterministic; graph-capturable once compiled (the first call per variant). + """ + + def __call__(self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject): + mixed, inject_out, _, _ = launch( + part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject + ) + return mixed, inject_out + + def reference(self, *args, **kwargs): + """The Triton entry's float64 reference.""" + return _triton_mix.reference(*args, **kwargs) + + def is_valid(self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject) -> None: + _triton_mix.is_valid(part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject) + _check(part, ss, up, hc_out, norm_weight, hc_count) + + +qwen4_exp_hc_decode_mix_cute = Qwen4ExpHcDecodeMixCute() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_tc_first.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_tc_first.py new file mode 100644 index 000000000000..e476bd967313 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_tc_first.py @@ -0,0 +1,670 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Hyper-connection mixer first half at decode row counts (2..128 rows) as one +tcgen05 split-K GEMM: the first projection `normed @ [down; block_inject]^T` +with the SiLU and injection activations in its epilogue, so neither the fp32 +logits nor a split-K reduce kernel nor `qwen4_exp_hc_lowrank_act` exists. + +Tier 6 (CuTe DSL written for this target). One CTA computes one +`[128 rows, BN columns]` partial over `K / KS` of the 10240-wide stream: a TMA +warp streams the normed rows (rows past `T` zero-filled by TMA, never read) +and the weight rows (only the `n_cols` real ones; the load-time zero padding +to 384 rows is never read), one MMA warp accumulates in tensor memory, four +epilogue warps move the fp32 partial into shared memory. The `KS` CTAs of one +column tile form a thread-block cluster; after a cluster barrier each CTA sums +a slice of the tile's `KS` partials out of its peers' shared memory (DSMEM) in +a fixed order (deterministic) and writes the activations. + +Launched with programmatic dependent launch: the weight stages stream before +the grid-dependency wait (the producer, `qwen4_exp_hc_inject_norm`, releases +dependents at its start), and this kernel releases its own consumer +(`qwen4_exp_hc_up_gate_mean(_fp4)`) at once so that it prefetches its weight. +""" + +from __future__ import annotations + +import os + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import torch +import torch.nn.functional as F +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.cute_dsl_kernels.blackwell.custom_pipeline import ( + PipelineTmaUmma, + PipelineUmmaAsync, +) +from tensorrt_llm._torch.cute_dsl_kernels.blackwell.utils import make_ptr + +#: Largest row count this entry is certified for: one 128-row MMA tile. +MAX_ROWS = 128 + +#: Programmatic dependent launch (sm_90+), off with TRTLLM_ENABLE_PDL=0. Without it +#: the grid-dependency wait and the dependents trigger are no-ops. +_PDL = os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" + + +def _asm(ret, args, text, constraints, loc, ip, side_effects=False): + return llvm.inline_asm( + ret, + [a.ir_value(loc=loc, ip=ip) for a in args], + text, + constraints, + has_side_effects=side_effects, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _mapa(addr, rank, *, loc=None, ip=None): + """This CTA's shared-memory address `addr` in cluster CTA `rank`'s window.""" + return cutlass.Int32( + _asm( + T.i32(), + [cutlass.Int32(addr), cutlass.Int32(rank)], + "mapa.shared::cluster.u32 $0, $1, $2;", + "=r,r,r", + loc, + ip, + ) + ) + + +@dsl_user_op +def _ld_dsmem_f32(addr, *, loc=None, ip=None): + """`ld.shared::cluster.f32`; volatile so it stays behind the cluster barrier.""" + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Int32(addr)], + "ld.shared::cluster.f32 $0, [$1];", + "=f,r", + loc, + ip, + side_effects=True, + ) + ) + + +class HcDecodeTcFirstKernel: + """`low[t, n] = silu(L[t, n] / C)` for `n < R` and, with `with_inject`, + `inject[t, j] = 2 * sigmoid(L[t, R + j] / C)`, where `L = A @ B^T`, + `A = normed [T, K]`, `B = weight rows [0, n_cols)` of `[>= n_cols, K]`. + + Compile-time: `bn` (MMA N, a multiple of 16), `ks` (K split = cluster + size), `hc`, `rank`, `n_cols`, `k`, `with_inject`. Runtime: the row count + (<= 128) and the pointers. + """ + + def __init__(self, bn, ks, hc, rank, n_cols, k, with_inject, epi_cols=16): + assert bn % epi_cols == 0 and bn % 16 == 0 and 16 <= bn <= 256 + assert k % (64 * ks) == 0, "K must split into 64-wide tiles per CTA" + self.bn, self.ks, self.hc, self.rank, self.n_cols, self.k = bn, ks, hc, rank, n_cols, k + self.with_inject = with_inject + self.n_tiles = (n_cols + bn - 1) // bn + self.acc_dtype = cutlass.Float32 + self.io_dtype = cutlass.BFloat16 + self.mma_tiler = (128, bn, 1) + self.epi_cols = epi_cols + self.cta_group = tcgen05.CtaGroup.ONE + # The MMA itself is single-CTA: no multicast; the launch cluster is along K only. + self.cluster_shape_mn = (1, 1) + self.epilogue_warp_id = (0, 1, 2, 3) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.threads_per_cta = 32 * 6 + self.epilog_sync_bar_id = 1 + self.tmem_alloc_sync_bar_id = 2 + self.num_acc_stage = 1 + # Flat elements of one column tile each CTA reduces, at most, per epilogue thread. + per_cta = (128 * bn + ks - 1) // ks + self.red_iters = (per_cta + self.threads_per_cta - 1) // self.threads_per_cta + + def _setup(self, a_major, b_major): + tiled_mma = utils.sm100.make_trivial_tiled_mma( + self.io_dtype, a_major, b_major, self.acc_dtype, self.cta_group, self.mma_tiler[:2] + ) + inst_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + self.mma_tiler = (self.mma_tiler[0], self.mma_tiler[1], inst_k * 4) + self.cta_tile_shape_mnk = self.mma_tiler + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), (tiled_mma.thr_id.shape,) + ) + self.epi_tile = (self.mma_tiler[0], self.epi_cols) + a_one = utils.sm100.make_smem_layout_a(tiled_mma, self.mma_tiler, self.io_dtype, 1) + b_one = utils.sm100.make_smem_layout_b(tiled_mma, self.mma_tiler, self.io_dtype, 1) + ab_bytes = cute.size_in_bytes(self.io_dtype, a_one) + cute.size_in_bytes( + self.io_dtype, b_one + ) + smem_capacity = utils.get_smem_capacity_in_bytes() + self.k_tiles = self.k // self.mma_tiler[2] // self.ks + # The fp32 partial, one padded row per accumulator row (conflict-free row stores). + self.p_stride = self.bn + 1 + p_bytes = 128 * self.p_stride * 4 + self.num_ab_stage = max(2, min(self.k_tiles, (smem_capacity - 2048 - p_bytes) // ab_bytes)) + self.n_pre = min(self.num_ab_stage, self.k_tiles) + self.a_smem_layout_staged = utils.sm100.make_smem_layout_a( + tiled_mma, self.mma_tiler, self.io_dtype, self.num_ab_stage + ) + self.b_smem_layout_staged = utils.sm100.make_smem_layout_b( + tiled_mma, self.mma_tiler, self.io_dtype, self.num_ab_stage + ) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + acc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) + self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(acc_fake, arch="sm_100") + return tiled_mma + + @cute.jit + def __call__( + self, + m: cutlass.Int32, + a_ptr: cute.Pointer, + b_ptr: cute.Pointer, + low_t: cute.Tensor, + inj_ptr: cute.Pointer, + stream: cuda.CUstream, + ): + k, ks = self.k, self.ks + a = cute.make_tensor(a_ptr, cute.make_layout((m, k, 1), stride=(k, 1, 0))) + b = cute.make_tensor(b_ptr, cute.make_layout((self.n_cols, k, 1), stride=(k, 1, 0))) + # `low` arrives as a tensor (the TVM-FFI entry takes its stream from a GPU tensor). + low = cute.make_tensor( + low_t.iterator, cute.make_layout((m, self.rank), stride=(self.rank, 1)) + ) + inj = cute.make_tensor(inj_ptr, cute.make_layout((m, self.hc), stride=(self.hc, 1))) + k_major = utils.LayoutEnum.ROW_MAJOR.mma_major_mode() + tiled_mma = self._setup(k_major, k_major) + + a_op = utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id) + a_smem = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_a, tma_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, a, a_smem, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape + ) + b_op = utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id) + b_smem = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, b, b_smem, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape + ) + self.num_tma_load_bytes = cute.size_in_bytes(self.io_dtype, a_smem) + cute.size_in_bytes( + self.io_dtype, b_smem + ) + self.kernel( + tiled_mma, + tma_atom_a, + tma_a, + tma_atom_b, + tma_b, + low, + inj, + m, + self.cluster_layout_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.epi_tile, + ).launch( + grid=[ks, self.n_tiles, 1], + block=[self.threads_per_cta, 1, 1], + cluster=[ks, 1, 1], + stream=stream, + use_pdl=_PDL, + ) + + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + mLow: cute.Tensor, + mInj: cute.Tensor, + m: cutlass.Int32, + cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + epi_tile: cute.Tile, + ): + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + # The consumer (the up projection, launched with programmatic dependent launch) + # may start at once: it streams its own static weight, then waits for this grid. + cute.arch.griddepcontrol_launch_dependents() + if warp_idx == self.tma_warp_id: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, _ = cute.arch.block_idx() + kr = bidx # K slice == rank in the (KS, 1, 1) cluster + nt = bidy # column tile + + @cute.struct + class SharedStorage: + ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] + tmem_dealloc_mbar: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + ab_producer, ab_consumer = PipelineTmaUmma.create( + barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + ).make_participants() + acc_pipeline = PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len(self.epilogue_warp_id) + ), + cta_layout_vmnk=cluster_layout_vmnk, + ) + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_alloc_sync_bar_id, + num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), + ) + tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.epilogue_warp_id[0], + is_two_cta=False, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr, + ) + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + + sA = smem.allocate_tensor( + element_type=self.io_dtype, + layout=a_smem_layout_staged.outer, + byte_alignment=128, + swizzle=a_smem_layout_staged.inner, + ) + sB = smem.allocate_tensor( + element_type=self.io_dtype, + layout=b_smem_layout_staged.outer, + byte_alignment=128, + swizzle=b_smem_layout_staged.inner, + ) + sP = smem.allocate_tensor( + element_type=self.acc_dtype, + layout=cute.make_layout((128, self.bn), stride=(self.p_stride, 1)), + byte_alignment=16, + ) + + gA_mkl = cute.local_tile( + mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + k0 = kr * self.k_tiles + + thr_mma = tiled_mma.get_slice(0) + tCgA = thr_mma.partition_A(gA_mkl) + tCgB = thr_mma.partition_B(gB_nkl) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + 0, + cute.make_layout(1), + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + 0, + cute.make_layout(1), + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + tCrA = tiled_mma.make_fragment_A(sA) + tCrB = tiled_mma.make_fragment_B(sB) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) + + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + if warp_idx == self.tma_warp_id: + tAgA_slice = tAgA[(None, 0, None, 0)] + tBgB_slice = tBgB[(None, nt, None, 0)] + ab_producer.reset() + # The weight does not depend on the previous kernel: the first stages' weight + # tiles stream before the grid-dependency wait (programmatic dependent + # launch), the normed rows -- the previous kernel's output -- after it. + pre = [] + for kk in cutlass.range_constexpr(self.n_pre): + handle = ab_producer.acquire_and_advance() + cute.copy( + tma_atom_b, + tBgB_slice[(None, k0 + kk)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + ) + pre.append(handle) + cute.arch.griddepcontrol_wait() + for kk in cutlass.range_constexpr(self.n_pre): + handle = pre[kk] + cute.copy( + tma_atom_a, + tAgA_slice[(None, k0 + kk)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + ) + for k_tile in cutlass.range(self.n_pre, self.k_tiles, 1, unroll=1): + handle = ab_producer.acquire_and_advance() + cute.copy( + tma_atom_a, + tAgA_slice[(None, k0 + k_tile)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, k0 + k_tile)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + ) + ab_producer.tail() + + if warp_idx == self.mma_warp_id: + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage + ) + tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] + ab_consumer.reset() + peek = ab_consumer.try_wait() + acc_pipeline.producer_acquire(acc_producer_state) + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for k_tile in cutlass.range(0, self.k_tiles, 1, unroll=1): + handle = ab_consumer.wait_and_advance(peek) + num_kblocks = cute.size(tCrA, mode=[2]) + for kb in cutlass.range(num_kblocks, unroll_full=True): + kcrd = (None, None, kb, handle.index) + cute.gemm(tiled_mma, tCtAcc, tCrA[kcrd], tCrB[kcrd], tCtAcc) + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + handle.release() + peek = cutlass.Boolean(1) + if handle.count + 1 < self.k_tiles: + peek = ab_consumer.try_wait() + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + acc_pipeline.producer_tail(acc_producer_state) + + if warp_idx < self.mma_warp_id: + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + copy_atom_t2r = utils.sm100.get_tmem_load_op( + self.cta_tile_shape_mnk, + utils.LayoutEnum.ROW_MAJOR, + self.acc_dtype, + self.acc_dtype, + epi_tile, + False, + ) + tAcc_epi = cute.flat_divide(tCtAcc_base[((None, None), 0, 0, None)], epi_tile) + tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)]) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + tTR_tAcc_base = thr_copy_t2r.partition_S(tAcc_epi) + tTR_sP = thr_copy_t2r.partition_D(cute.flat_divide(sP, epi_tile)) + frag_shape = tTR_sP[(None, None, None, 0, 0)].shape + tTR_rAcc = cute.make_rmem_tensor(frag_shape, self.acc_dtype) + tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_consumer_state.index)] + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + acc_pipeline.consumer_wait(acc_consumer_state) + for sub in cutlass.range_constexpr(self.bn // self.epi_cols): + cute.copy(tiled_copy_t2r, tTR_tAcc[(None, None, None, sub)], tTR_rAcc) + cute.autovec_copy(tTR_rAcc, tTR_sP[(None, None, None, 0, sub)]) + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + cute.arch.barrier( + barrier_id=self.epilog_sync_bar_id, + number_of_threads=32 * len(self.epilogue_warp_id), + ) + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + # Release / acquire at cluster scope, every thread of every CTA: after it, each + # CTA reads all `KS` partials of its column tile out of its peers' shared memory. + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + + # This CTA's share of the tile's `m * BN` outputs, all 6 warps. Every load is + # issued before any sum (indices clamped in bounds, so no branch separates + # them); the `KS` partials are summed in slice order (deterministic). + total = m * self.bn + per = (total + self.ks - 1) // self.ks + base = kr * per + nthr = self.threads_per_cta + vals = cute.make_rmem_tensor((self.red_iters * self.ks,), self.acc_dtype) + p_base = cutlass.Int32(sP.iterator.toint()) + peers = [_mapa(p_base, j) for j in range(self.ks)] + for it in cutlass.range_constexpr(self.red_iters): + fs = cutlass.min(base + tidx + it * nthr, total - 1) + row = fs // self.bn + col = fs - row * self.bn + off = (row * self.p_stride + col) * 4 + for j in cutlass.range_constexpr(self.ks): + vals[it * self.ks + j] = _ld_dsmem_f32(peers[j] + off) + inv_hc = cutlass.Float32(1.0 / self.hc) + for it in cutlass.range_constexpr(self.red_iters): + local = tidx + it * nthr + f = base + local + if local < per and f < total: + row = f // self.bn + col = f - row * self.bn + n = nt * self.bn + col + acc = vals[it * self.ks] + for j in cutlass.range_constexpr(1, self.ks): + acc = acc + vals[it * self.ks + j] + z = acc * inv_hc + sg = cutlass.Float32(1.0) / (cutlass.Float32(1.0) + cute.math.exp(-z)) + if n < self.rank: + mLow[(row, n)] = (z * sg).to(self.io_dtype) + if cutlass.const_expr(self.with_inject): + if n >= self.rank and n < self.rank + self.hc: + mInj[(row, n - self.rank)] = (cutlass.Float32(2.0) * sg).to(self.io_dtype) + + # No CTA may exit (releasing its shared memory) while a peer still reads it. + cute.arch.cluster_arrive_relaxed() + cute.arch.cluster_wait() + + +_COMPILED: dict = {} + + +def _launcher(bn, ks, hc, rank, n_cols, k, with_inject, ptrs): + """Compile once per static configuration; the row count is a runtime argument + and the TVM-FFI entry launches on the caller's current stream (graph-capturable). + `ptrs` is `(normed, weight, low tensor, inject)`.""" + key = (bn, ks, hc, rank, n_cols, k, with_inject) + fn = _COMPILED.get(key) + if fn is None: + kern = HcDecodeTcFirstKernel(bn, ks, hc, rank, n_cols, k, with_inject) + gmem = cute.AddressSpace.gmem + cptrs = [ + make_ptr(cutlass.BFloat16, ptrs[0], gmem, assumed_align=16), + make_ptr(cutlass.BFloat16, ptrs[1], gmem, assumed_align=16), + from_dlpack(ptrs[2], assumed_align=16).mark_layout_dynamic(leading_dim=1), + make_ptr(cutlass.BFloat16, ptrs[3], gmem, assumed_align=2), + ] + fn = cute.compile( + kern, + MAX_ROWS, + *cptrs, + stream=cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--opt-level 2 --enable-tvm-ffi", + ) + _COMPILED[key] = fn + return fn + + +def hc_decode_tc_first(normed, first, rank, hc_count, with_inject, bn, ks): + """Launch with an explicit tile (`bn` columns, `ks`-way K split). `first` is + the `[K, >= n_cols]` transposed view of the row-major weight, as + `cublas_mm(normed, first)` takes it. Returns `(low, inject)`.""" + rows, k = normed.shape + n_cols = rank + (hc_count if with_inject else 0) + dev, dt = normed.device, normed.dtype + low = torch.empty((rows, rank), dtype=dt, device=dev) + inject = torch.empty((rows, hc_count), dtype=dt, device=dev) if with_inject else None + if rows > 0: + ptrs = ( + normed.data_ptr(), + first.data_ptr(), + low, + (inject if with_inject else low).data_ptr(), + ) + _launcher(bn, ks, hc_count, rank, n_cols, k, with_inject, ptrs)(rows, *ptrs) + return low, inject + + +#: Tile: `_BN` weight rows per CTA (7 tiles cover the 324 real rows of 336) and +#: a `_KS`-way K split, one cluster per column tile. +_BN = 48 +_KS = 16 + + +class Qwen4ExpHcDecodeTcFirst(OpWrapper): + """With `R = rank`, `C = hc_count`, `normed` the `[T, K]` grouped-normed + stream and `first` the `[K, >= R (+ C)]` transposed view of the load-time + `[down; block_inject; 0]` weight, `L = normed @ first` in fp32:: + + low = silu(L[:, :R] / C) # [T, R], normed's dtype + inject = 2 * sigmoid(L[:, R : R + C] / C) # [T, C], only with_inject + + The projection accumulates in fp32 (tensor memory, then a fixed-order sum + of the K-split partials); each activation is fp32 with one rounding. + Returns fresh `(low, inject)`; `inject` is None exactly when `with_inject` + is False. Fusion boundary: what `cublas_mm(normed, first, None, fp32)` + + `qwen4_exp_hc_lowrank_act` compute (cuBLAS's split-K reduce included). + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: rows, inject (True), hc_count (4), rank (320), hidden (2560), n_rows + # (384: the padded weight), dtype bf16 only. + CELLS: tuple[Cell, ...] = tuple( + Cell(why=f"checkpoint 4x2560, 324 cols of 384, inject, {why}", spec=dict(rows=t)) + for t, why in ( + (2, "decode 2 rows"), + (5, "decode 5 rows"), + (16, "decode 16 rows"), + (33, "decode 33 rows"), + (64, "decode 64 rows"), + (100, "decode 100 rows"), + (128, "decode 128 rows (max concurrency, MAX_ROWS)"), + ) + ) + ( + Cell(why="stack-final mixer: 320 cols, no inject, 2 rows", spec=dict(rows=2, inject=False)), + Cell( + why="stack-final mixer: 320 cols, no inject, 128 rows", + spec=dict(rows=128, inject=False), + ), + Cell(why="unpadded weight [324, K], 8 rows", spec=dict(rows=8, n_rows=324)), + Cell(why="hc_count 2 at hidden 2560, 24 rows", spec=dict(rows=24, hc_count=2)), + Cell(why="1 row (below the dispatch range)", spec=dict(rows=1)), + Cell(why="zero rows", spec=dict(rows=0)), + ) + + note = """ + Tier 6 (CuTe DSL written for this target): tiers 1-5 have no GEMM whose + epilogue applies an activation after a split-K reduction; cuBLAS at + N = 384 runs split-K plus its own reduce kernel and writes fp32 logits + that `qwen4_exp_hc_lowrank_act` reads back (three launches), and the + Triton `tl.dot` K-loop designs measured 36-60 us at 128 rows. + Metadata consumed: none. Call sites: `DecodeTarget`'s trunk + hyper-connection mixers from `_HC_TC_FIRST_MIN_ROWS` to + `_HC_TC_FIRST_MAX_ROWS` (= MAX_ROWS) rows, 97 per forward, between + `qwen4_exp_hc_inject_norm` (or the PLE tail) and + `qwen4_exp_hc_up_gate_mean(_fp4)`. Preconditions (bf16): `normed` + `[T, K]` contiguous, T <= 128; `first` the `.t()` view of a row-major + `[N, K]` weight with N >= R (+ C); 16-byte aligned; K a multiple of + 64 * KS. Only the R (+ C) real weight rows are read (TMA zero-fills + the tile past them), never the load-time padding. + Launch: grid (KS, cdiv(R + C, BN)) = (16, 7) CTAs of 6 warps, cluster + (KS, 1, 1), one CTA per SM (~227 KB shared memory: every K tile of the + slice in flight plus the padded fp32 partial). The partials are reduced + through distributed shared memory between two cluster barriers; no + global workspace. Deterministic (fixed K-slice order). Programmatic + dependent launch (TRTLLM_ENABLE_PDL=0 turns it off): weight stages + before `griddepcontrol.wait`, `normed` after it; dependents released + at kernel start. Graph-capturable (current stream, no host sync, no + persistent state). Compiles once per (BN, KS, C, R, N, K, inject). + Measured on GB300 (graph replay, 24 cold weights, `microbench/` of the + candidate): 6.4 us at 128 rows against 8.6 us for cuBLAS split-K + + reduce + lowrank_act; tile sweep BN 32 / 48 / 64 / 112 x KS 8 / 10 / + 16, 48 x 16 fastest at every row count. + """ + + def __call__( + self, + normed: torch.Tensor, + first: torch.Tensor, + rank: int, + hc_count: int, + with_inject: bool, + ): + return hc_decode_tc_first(normed, first, rank, hc_count, with_inject, _BN, _KS) + + def reference( + self, + normed: torch.Tensor, + first: torch.Tensor, + rank: int, + hc_count: int, + with_inject: bool, + ): + """float64 projection and activations, one rounding per output.""" + n_cols = rank + (hc_count if with_inject else 0) + lg = normed.double() @ first[:, :n_cols].double() + low = F.silu(lg[:, :rank] / hc_count).to(normed.dtype) + inject = None + if with_inject: + inject = (2.0 * torch.sigmoid(lg[:, rank:n_cols] / hc_count)).to(normed.dtype) + return low, inject + + def is_valid( + self, + normed: torch.Tensor, + first: torch.Tensor, + rank: int, + hc_count: int, + with_inject: bool, + ) -> None: + rows, k = normed.shape + n_cols = rank + (hc_count if with_inject else 0) + assert normed.dtype == first.dtype == torch.bfloat16, "bf16 operands" + assert rows <= MAX_ROWS, f"at most {MAX_ROWS} rows (one MMA tile)" + assert normed.stride() == (k, 1), "normed contiguous [T, K]" + assert first.shape[0] == k and first.shape[1] >= n_cols, "first [K, >= n_cols]" + assert first.stride() == (1, k), "first must be the .t() view of a row-major [N, K] weight" + assert k % (64 * _KS) == 0, f"K must be a multiple of {64 * _KS}" + for t in (normed, first): + assert t.data_ptr() % 16 == 0, "16-byte aligned operands" + + +qwen4_exp_hc_decode_tc_first = Qwen4ExpHcDecodeTcFirst() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_gate_mean.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_gate_mean.py new file mode 100644 index 000000000000..0aaa1bd5bcd5 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_gate_mean.py @@ -0,0 +1,125 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Hyper-connection mixer epilogue: the sigmoid gate applied to the normed +stream and the mean over streams, from the up projection's logits, in one +Triton kernel. + +Tier 5 (Triton written for this target). Used where the up projection runs on +cuBLAS (decode below `_HC_FUSED_UP_MIN_ROWS` rows): `qwen4_exp_hc_lowrank_act` -> +`cublas_mm(low, up^T)` -> this. Elsewhere `qwen4_exp_hc_up_gate_mean` computes the +up projection and this epilogue in one kernel. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +@triton.jit +def _hc_gate_mean_kernel( + gate_ptr, + normed_ptr, + mixed_ptr, + stride_gate, + stride_normed, + stride_mixed, + HIDDEN: tl.constexpr, + HC: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + h = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + mask = h < HIDDEN + acc = tl.zeros([BLOCK], dtype=tl.float32) + for s in tl.static_range(HC): + g = tl.load(gate_ptr + row * stride_gate + s * HIDDEN + h, mask=mask, other=0.0) + x = tl.load(normed_ptr + row * stride_normed + s * HIDDEN + h, mask=mask, other=0.0) + acc += tl.sigmoid(g.to(tl.float32)) * x.to(tl.float32) + tl.store( + mixed_ptr + row * stride_mixed + h, + (acc * (1.0 / HC)).to(mixed_ptr.dtype.element_ty), + mask=mask, + ) + + +class Qwen4ExpHcGateMean(OpWrapper): + """With `C = hc_count`, `H = width // C`, `gate_logits` and `normed` both + `[T, C*H]` stream-major:: + + mixed[t, i] = mean over s of sigmoid(gate_logits[t, s*H + i]) * normed[t, s*H + i] + + fp32 inside, one rounding. Returns a fresh `[T, H]` tensor in `normed`'s + dtype. Fusion boundary: the sigmoid, the product and the stream mean; the + up projection producing `gate_logits` is the caller's (`cublas_mm`). + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: rows, hc_count (4), hidden (2560), dtype (bf16). + CELLS: tuple[Cell, ...] = tuple( + Cell(why=f"checkpoint 4x2560 bf16, {why}", spec=dict(rows=t)) + for t, why in ( + (1, "1 row"), + (128, "128 rows"), + (2048, "prefill 2048"), + (8192, "prefill 8192"), + ) + ) + ( + Cell(why="hc_count 2, 64 rows", spec=dict(rows=64, hc_count=2)), + Cell(why="hc_count 8, 64 rows", spec=dict(rows=64, hc_count=8)), + Cell(why="hidden 128 (not the checkpoint), 48 rows", spec=dict(rows=48, hidden=128)), + Cell(why="float16 operands, 64 rows", spec=dict(rows=64, dtype=torch.float16)), + Cell(why="zero rows", spec=dict(rows=0)), + ) + + note = """ + Metadata consumed: none. Call sites: the trunk's hyper-connection sites in + `DecodeTarget` below `_HC_FUSED_UP_MIN_ROWS` rows, after `cublas_mm(low, up^T)`. Preconditions: both + operands 2-D `[T, C*H]` with unit column stride. Graph-capturable. + Measured (GB300, graph): the three-launch cuBLAS route this ends is 2.7x + faster than `qwen4_exp_hc_gate_mix` at 512..8192 rows (106 vs 287 us at + 8192) and 1.7x at 128. + """ + + def __call__( + self, gate_logits: torch.Tensor, normed: torch.Tensor, hc_count: int + ) -> torch.Tensor: + rows, width = normed.shape + hidden = width // hc_count + mixed = torch.empty((rows, hidden), dtype=normed.dtype, device=normed.device) + if rows > 0: + block = min(1024, triton.next_power_of_2(hidden)) + _hc_gate_mean_kernel[(rows, triton.cdiv(hidden, block))]( + gate_logits, + normed, + mixed, + gate_logits.stride(0), + normed.stride(0), + mixed.stride(0), + HIDDEN=hidden, + HC=hc_count, + BLOCK=block, + num_warps=4, + ) + return mixed + + def reference( + self, gate_logits: torch.Tensor, normed: torch.Tensor, hc_count: int + ) -> torch.Tensor: + """float64, one rounding.""" + hidden = normed.shape[1] // hc_count + gate = torch.sigmoid(gate_logits.double()).unflatten(-1, (hc_count, hidden)) + x = normed.double().unflatten(-1, (hc_count, hidden)) + return (gate * x).mean(dim=-2).to(normed.dtype) + + def is_valid(self, gate_logits: torch.Tensor, normed: torch.Tensor, hc_count: int) -> None: + assert gate_logits.shape == normed.shape, "gate_logits and normed must match" + assert gate_logits.stride(1) == 1 and normed.stride(1) == 1, "unit column stride" + assert normed.shape[1] % hc_count == 0 + + +qwen4_exp_hc_gate_mean = Qwen4ExpHcGateMean() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_inject_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_inject_norm.py new file mode 100644 index 000000000000..347942e58938 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_inject_norm.py @@ -0,0 +1,1088 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Hyper-connection post-block injection fused with the next site's grouped +gemma RMS norm: one Triton row kernel writes the updated multi-stream residual +and its normalized copy in the same pass. After a MoE block it also forms the +block output itself, from the routed experts' output, the shared expert's +output and its gate logit (the shared-expert combine). + +Tier 5 (Triton written for this target). No ``torch.ops.trtllm.*`` op or +FlashInfer callable accumulates a block output into ``hc_count`` streams with +per-stream gates and normalizes each stream with its own RMS in one launch; the +two-launch composition (``qwen4_exp_post_inject`` + four +``flashinfer_gemma_rmsnorm`` calls and a concat) is what this replaces, and +with ``shared`` / ``gate_logits`` also the ``fused_sigmoid_gate_mul_add`` +launch before it -- and with ``topk_weights`` / ``topk_rows`` the MoE runner's +finalize launch as well. + +The decode MoE sites (the finalize fold below `_FEW_WARPS_MIN_ROWS` rows) run a +CuTe DSL kernel instead (Tier 6, written for this target), bitwise the Triton +kernel's outputs there. +""" + +from __future__ import annotations + +import os +import struct +from typing import Optional, Tuple + +import cutlass +import cutlass.cute as cute +import torch +import triton +import triton.language as tl +from cutlass._mlir.dialects import llvm +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.cute_dsl_kernels.blackwell.utils import make_ptr + +#: Trigger dependents at kernel start (`griddepcontrol.launch_dependents`); a no-op +#: unless the next kernel is launched with programmatic dependent launch. +_PDL_TRIGGER = os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" + +#: Rows from which a wide (`hidden >= 2048`) launch runs two warps a program +#: instead of eight: prefill-sized grids are bandwidth-bound and run faster +#: with fewer, longer threads; decode-sized ones keep eight. +_FEW_WARPS_MIN_ROWS = 512 + + +@triton.jit +def _hc_inject_norm_kernel( + hc_ptr, + block_ptr, + inject_ptr, + shared_ptr, + gate_ptr, + topk_w_ptr, + topk_row_ptr, + weight_ptr, + hc_out_ptr, + normed_ptr, + stride_hc, + stride_block, + stride_inject, + stride_shared, + stride_gate, + stride_topk_w, + stride_topk_row, + stride_out, + stride_normed, + eps, + HC_COUNT: tl.constexpr, + HIDDEN: tl.constexpr, + BLOCK: tl.constexpr, + HAS_INJECT: tl.constexpr, + HAS_SHARED: tl.constexpr, + TOPK: tl.constexpr, + STREAMS_PER_PROGRAM: tl.constexpr, + PDL_TRIGGER: tl.constexpr, +): + if PDL_TRIGGER: + # Release a programmatic-dependent-launch consumer (the decode mixer's + # first projection) at once: it streams its static weights while this + # kernel runs and waits for the grid to finish before reading `normed`. + tl.extra.cuda.gdc_launch_dependents() + # One program per `STREAMS_PER_PROGRAM` streams of a row; a stream's `HIDDEN` + # columns are one RMS group. A row's programs are adjacent in launch order, so + # the row's block output (and the combine's operands) come from HBM once and + # from L2 for the others; with several streams a program forms it once for them. + pid = tl.program_id(0).to(tl.int64) + row = pid // (HC_COUNT // STREAMS_PER_PROGRAM) + first_stream = ((pid % (HC_COUNT // STREAMS_PER_PROGRAM)) * STREAMS_PER_PROGRAM).to(tl.int32) + cols = tl.arange(0, BLOCK) + mask = cols < HIDDEN + if HAS_INJECT: + if TOPK > 0: + # The routed experts' finalize: this row's TOPK permuted FC2 rows weighted + # by their routing weights, fp32 FMAs in slot order from zero (a slot mapped + # to -1 is skipped), rounded once to the rows' dtype. + b = tl.zeros([BLOCK], dtype=tl.float32) + for k in tl.static_range(TOPK): + prow = tl.load(topk_row_ptr + row * stride_topk_row + k) + wk = tl.load(topk_w_ptr + row * stride_topk_w + k).to(tl.float32) + r = tl.load( + block_ptr + tl.maximum(prow, 0).to(tl.int64) * stride_block + cols, + mask=mask, + other=0.0, + ).to(tl.float32) + b = tl.where(prow >= 0, tl.fma(r, wk, b), b) + b = b.to(block_ptr.dtype.element_ty).to(tl.float32) + else: + b = tl.load(block_ptr + row * stride_block + cols, mask=mask, other=0.0).to(tl.float32) + if HAS_SHARED: + # The MoE block output, `routed + sigmoid(gate) * shared` in fp32, rounded + # once to the routed output's dtype -- the value that block used to store. + sg = tl.sigmoid(tl.load(gate_ptr + row * stride_gate).to(tl.float32)) + sh = tl.load(shared_ptr + row * stride_shared + cols, mask=mask, other=0.0) + b = (b + sg * sh.to(tl.float32)).to(block_ptr.dtype.element_ty).to(tl.float32) + for j in tl.static_range(STREAMS_PER_PROGRAM): + stream = first_stream + j + offs = stream * HIDDEN + cols + x = tl.load(hc_ptr + row * stride_hc + offs, mask=mask, other=0.0).to(tl.float32) + if HAS_INJECT: + g = tl.load(inject_ptr + row * stride_inject + stream).to(tl.float32) + # Rounded once on the store, like `addcmul`; the norm reads the rounded + # residual, which is the value every later site sees. + xr = (x + b * g).to(hc_out_ptr.dtype.element_ty) + tl.store(hc_out_ptr + row * stride_out + offs, xr, mask=mask) + x = xr.to(tl.float32) + rstd = tl.rsqrt(tl.sum(x * x, axis=0) / HIDDEN + eps) + w = tl.load(weight_ptr + offs, mask=mask, other=0.0).to(tl.float32) + y = x * rstd * (1.0 + w) + tl.store( + normed_ptr + row * stride_normed + offs, y.to(normed_ptr.dtype.element_ty), mask=mask + ) + + +# -------------------------------------------------------------------------------------- +# The decode MoE sites in CuTe DSL. There the Triton kernel runs two programs a row, each +# gathering all `k` permuted FC2 rows and normalizing its two streams one after the other +# (6.5 us at 128 rows in the con128 decode step on GB300). Here the row's two CTAs are a +# cluster: CTA j forms the block output's columns [j * H / 2, (j + 1) * H / 2) -- the +# finalize's slot-order FMA chain and the combine, the Triton kernel's instructions -- and +# stores them into both CTAs' shared memory; after a cluster barrier each CTA reads the +# whole block row locally and accumulates and normalizes its two streams together. A CTA +# is 256 threads in the Triton kernel's layout at a 4096 block and eight warps (thread t: +# columns [8t, 8t + 8) and, below the hidden, [2048 + 8t, 2048 + 8t + 8)), and the sum of +# squares is its reduction tree: `hc_out` and `normed` are bitwise the Triton kernel's. +# -------------------------------------------------------------------------------------- + +_CUTE_THREADS = 256 +_CUTE_WARPS = _CUTE_THREADS // 32 +_CUTE_VEC = 8 +#: The Triton kernel's block at these hiddens is twice this: a thread's second chunk starts here. +_CUTE_SPLIT = _CUTE_THREADS * _CUTE_VEC +#: Streams a CTA, as the Triton fold's program; a row's `hc_count // 2` CTAs are one cluster. +_CUTE_SPP = 2 + + +def _iv(x, loc, ip): + return x.ir_value(loc=loc, ip=ip) + + +def _asm(ret, args, text, constraints, loc, ip, side_effects=False): + return llvm.inline_asm( + ret, + [_iv(a, loc, ip) for a in args], + text, + constraints, + has_side_effects=side_effects, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _asm4(args, text, constraints, loc, ip, side_effects=False): + ty = T.i32() + res = _asm( + llvm.StructType.get_literal([ty] * 4), args, text, constraints, loc, ip, side_effects + ) + return [cutlass.Int32(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip)) for i in range(4)] + + +@dsl_user_op +def _ldg_v4(addr, pred, *, loc=None, ip=None): + """16 bytes, or zeros where `pred` is 0 (nothing read).""" + return _asm4( + [cutlass.Int64(addr), cutlass.Int32(pred)], + "{ .reg .pred p; setp.ne.b32 p, $5, 0; mov.b32 $0, 0; mov.b32 $1, 0; mov.b32 $2, 0; " + "mov.b32 $3, 0; @p ld.global.v4.b32 {$0, $1, $2, $3}, [$4]; }", + "=r,=r,=r,=r,l,r", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _stg_v4(addr, pred, a, b, c, d, *, loc=None, ip=None): + """16 bytes where `pred` is not 0.""" + _asm( + T.i32(), + [cutlass.Int64(addr), cutlass.Int32(pred)] + [cutlass.Int32(v) for v in (a, b, c, d)], + "{ .reg .pred p; setp.ne.b32 p, $2, 0; @p st.global.v4.b32 [$1], {$3, $4, $5, $6}; " + "mov.u32 $0, 0; }", + "=r,l,r,r,r,r,r", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _mapa(addr, rank, *, loc=None, ip=None): + """This CTA's shared-memory address `addr` in cluster CTA `rank`'s window.""" + return cutlass.Int32( + _asm( + T.i32(), + [cutlass.Int32(addr), cutlass.Int32(rank)], + "mapa.shared::cluster.u32 $0, $1, $2;", + "=r,r,r", + loc, + ip, + ) + ) + + +@dsl_user_op +def _st_dsmem_v4(addr, pred, a, b, c, d, *, loc=None, ip=None): + """16 bytes into a cluster CTA's shared memory where `pred` is not 0.""" + _asm( + T.i32(), + [cutlass.Int32(addr), cutlass.Int32(pred)] + [cutlass.Int32(v) for v in (a, b, c, d)], + "{ .reg .pred p; setp.ne.b32 p, $2, 0; @p st.shared::cluster.v4.b32 [$1], " + "{$3, $4, $5, $6}; mov.u32 $0, 0; }", + "=r,r,r,r,r,r,r", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _lds_v4(addr, pred, *, loc=None, ip=None): + """16 bytes of this CTA's shared memory, or zeros where `pred` is 0.""" + return _asm4( + [cutlass.Int32(addr), cutlass.Int32(pred)], + "{ .reg .pred p; setp.ne.b32 p, $5, 0; mov.b32 $0, 0; mov.b32 $1, 0; mov.b32 $2, 0; " + "mov.b32 $3, 0; @p ld.shared.v4.b32 {$0, $1, $2, $3}, [$4]; }", + "=r,=r,=r,=r,r,r", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _bf16_lo(w, *, loc=None, ip=None): + """The low bf16 of a word as fp32 (exact).""" + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Int32(w)], + "{ .reg .b32 t; shl.b32 t, $1, 16; mov.b32 $0, t; }", + "=f,r", + loc, + ip, + ) + ) + + +@dsl_user_op +def _bf16_hi(w, *, loc=None, ip=None): + """The high bf16 of a word as fp32 (exact).""" + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Int32(w)], + "{ .reg .b32 t; and.b32 t, $1, 0xFFFF0000; mov.b32 $0, t; }", + "=f,r", + loc, + ip, + ) + ) + + +@dsl_user_op +def _bf16x2(lo, hi, *, loc=None, ip=None): + """Two fp32 values rounded (RN) into one bf16x2 word, `lo` low.""" + return cutlass.Int32( + _asm( + T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "cvt.rn.bf16x2.f32 $0, $1, $2;", + "=r,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _round_bf16(x, *, loc=None, ip=None): + """fp32 -> bf16 (RN) -> fp32.""" + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Float32(x)], + "{ .reg .b16 h; cvt.rn.bf16.f32 h, $1; cvt.f32.bf16 $0, h; }", + "=f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _fma(a, b, c, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _fma_if(slot, a, b, c, *, loc=None, ip=None): + """`slot >= 0 ? fma(a, b, c) : c` (the finalize skips a slot mapped to -1).""" + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Int32(slot), cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)], + "{ .reg .pred p; .reg .f32 t; fma.rn.f32 t, $2, $3, $4; setp.ge.s32 p, $1, 0; " + "selp.f32 $0, t, $4, p; }", + "=f,r,f,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _mul(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _add(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "add.rn.f32 $0, $1, $2;", + "=f,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _sigmoid(x, *, loc=None, ip=None): + """`tl.sigmoid` as the Triton kernel compiles it: `mul` by -log2(e), `ex2.approx`, `add` + 1, `div.full`.""" + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Float32(x)], + "{ .reg .f32 t; mul.f32 t, $1, 0fBFB8AA3B; ex2.approx.f32 t, t; " + "add.f32 t, t, 0f3F800000; div.full.f32 $0, 0f3F800000, t; }", + "=f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _rstd(ss, eps, hidden_hex: str, *, loc=None, ip=None): + """`rsqrt(ss / hidden + eps)` as the Triton kernel's `div.full`, `add`, `rsqrt.approx.ftz`: + the `add` has no rounding mode there, so ptxas fuses it with the division (a multiply by + the reciprocal) into one FFMA -- here the same.""" + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Float32(ss), cutlass.Float32(eps)], + "{ .reg .f32 t; div.full.f32 t, $1, " + hidden_hex + "; add.f32 t, t, $2; " + "rsqrt.approx.ftz.f32 $0, t; }", + "=f,f,f", + loc, + ip, + ) + ) + + +def _vals(words): + """Four bf16x2 words -> their eight values as fp32, in column order.""" + out = [] + for w in words: + out += [_bf16_lo(w), _bf16_hi(w)] + return out + + +class _FoldNormKernel: + """CTA (j, row): the row's streams 2j and 2j + 1; the row's CTAs one cluster.""" + + def __init__(self, hc_count: int, hidden: int, top_k: int, pdl_trigger: bool): + self.hc, self.hidden, self.topk, self.pdl = hc_count, hidden, top_k, pdl_trigger + self.ctas = hc_count // _CUTE_SPP + # The block-output columns a CTA forms, one 16-byte vector a thread. + self.per = hidden // self.ctas + self.n_fold = self.per // _CUTE_VEC + # Threads with a second chunk (columns [2048, hidden)). + self.n_b = (hidden - _CUTE_SPLIT) // _CUTE_VEC + self.hidden_hex = "0f%08X" % struct.unpack("> off)) + if tidx % 32 == 0: + red[s, warp] = v + cute.arch.sync_threads() + for s in cutlass.range_constexpr(_CUTE_SPP): + stream = rank * _CUTE_SPP + s + v = red[s, tidx % _CUTE_WARPS] + for off in cutlass.range_constexpr(3): + v = _add(v, cute.arch.shuffle_sync_bfly(v, offset=4 >> off)) + rstd = _rstd(v, eps, self.hidden_hex) + ra, rb = xrs[s] + wa_v, wb_v = _vals(wa[s]), _vals(wb[s]) + ya = [ + _mul(_mul(rstd, ra[e]), _add(wa_v[e], cutlass.Float32(1.0))) + for e in range(_CUTE_VEC) + ] + yb = [ + _mul(_mul(rstd, rb[e]), _add(wb_v[e], cutlass.Float32(1.0))) + for e in range(_CUTE_VEC) + ] + nrow = nrm_base + (row64 * width + stream * self.hidden) * 2 + _stg_v4(nrow + c_a * 2, one, *[_bf16x2(ya[2 * i], ya[2 * i + 1]) for i in range(4)]) + _stg_v4(nrow + c_b * 2, has_b, *[_bf16x2(yb[2 * i], yb[2 * i + 1]) for i in range(4)]) + + +_FOLD_NORM: dict = {} + + +def _fold_cute_ok( + hyper_input, + block_out, + inject, + norm_weight, + hc_count, + shared, + gate_logits, + topk_weights, + topk_rows, +) -> bool: + """The decode MoE site this kernel certifies (the Triton kernel's eight-warp form with + 16-byte vectors); anything else runs the Triton kernel.""" + rows, width = hyper_input.shape + hidden = width // hc_count + bf = torch.bfloat16 + return ( + 0 < rows < _FEW_WARPS_MIN_ROWS + and hc_count == 4 + and _CUTE_SPLIT < hidden <= 2 * _CUTE_SPLIT + and hidden % (_CUTE_SPP * _CUTE_VEC) == 0 + and hyper_input.dtype == block_out.dtype == shared.dtype == bf + and inject.dtype == gate_logits.dtype == topk_weights.dtype == norm_weight.dtype == bf + and topk_rows.dtype == torch.int32 + and all( + t.data_ptr() % 16 == 0 and t.stride(0) % 16 == 0 + for t in (hyper_input, block_out, shared) + ) + and norm_weight.data_ptr() % 16 == 0 + ) + + +def _fold_cute( + hyper_input, + block_out, + inject, + norm_weight, + eps, + hc_count, + shared, + gate_logits, + topk_weights, + topk_rows, + hc_out, + normed, +) -> None: + rows, width = hyper_input.shape + hidden = width // hc_count + top_k = topk_rows.shape[1] + ptrs = ( + hyper_input, + block_out.data_ptr(), + inject.data_ptr(), + shared.data_ptr(), + gate_logits.data_ptr(), + topk_weights.data_ptr(), + topk_rows.data_ptr(), + norm_weight.data_ptr(), + hc_out.data_ptr(), + normed.data_ptr(), + ) + strides = ( + hyper_input.stride(0), + block_out.stride(0), + inject.stride(0), + shared.stride(0), + gate_logits.stride(0), + topk_weights.stride(0), + topk_rows.stride(0), + ) + key = (hc_count, hidden, top_k, _PDL_TRIGGER) + fn = _FOLD_NORM.get(key) + if fn is None: + gm = cute.AddressSpace.gmem + bf = cutlass.BFloat16 + cargs = ( + from_dlpack(hyper_input, assumed_align=16).mark_layout_dynamic(leading_dim=1), + make_ptr(bf, ptrs[1], gm, assumed_align=16), + make_ptr(bf, ptrs[2], gm, assumed_align=2), + make_ptr(bf, ptrs[3], gm, assumed_align=16), + make_ptr(bf, ptrs[4], gm, assumed_align=2), + make_ptr(bf, ptrs[5], gm, assumed_align=2), + make_ptr(cutlass.Int32, ptrs[6], gm, assumed_align=4), + make_ptr(bf, ptrs[7], gm, assumed_align=16), + make_ptr(bf, ptrs[8], gm, assumed_align=16), + make_ptr(bf, ptrs[9], gm, assumed_align=16), + ) + fn = cute.compile( + _FoldNormKernel(hc_count, hidden, top_k, _PDL_TRIGGER), + *cargs, + *(cutlass.Int32(1),) * 8, + cutlass.Float32(1e-6), + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--opt-level 2 --enable-tvm-ffi", + ) + _FOLD_NORM[key] = fn + fn(*ptrs, rows, *strides, eps) + + +class Qwen4ExpHcInjectNorm(OpWrapper): + """With `C = hc_count`, `H = width // C`, every tensor 2-D `[T, ...]`: + + hc_out[t, s*H + i] = hyper_input[t, s*H + i] + block_out[t, i] * inject[t, s] + normed[t, s*H + i] = hc_out[t, s*H + i] * rstd(t, s) * (1 + norm_weight[s*H + i]) + rstd(t, s) = 1 / sqrt(mean_i(hc_out[t, s*H + i]^2) + eps) + + `hc_out` is rounded once to the activation dtype and `normed` is computed + from that rounded value (fp32 inside, one rounding on the store). Returns + `(hc_out, normed)`, both fresh `[T, C*H]` tensors in `hyper_input`'s dtype. + + `block_out=None, inject=None` is the norm alone: `hc_out` *is* + `hyper_input` (not a copy) and only `normed` is written -- the stack's first + site and the site after the PLE add. + + With `shared` and `gate_logits` (after a MoE block), `block_out` is the + routed experts' output and the block output injected above is + + block[t, i] = block_out[t, i] + sigmoid(gate_logits[t]) * shared[t, i] + + in fp32, rounded once to `block_out`'s dtype before the injection reads it: + bitwise what `fused_sigmoid_gate_mul_add` stores. `gate_logits` is one + pre-sigmoid scalar per row, `[T, 1]` or `[T]`; none of the inputs is written. + + With `topk_weights` / `topk_rows` as well (the MoE runner left unfinalized), + `block_out` is its permuted FC2 output `[P, H]` and the routed row is the + runner's finalize, formed here: `routed[t, i] = sum_k topk_weights[t, k] * + block_out[topk_rows[t, k], i]` as fp32 FMAs in slot order from zero (a slot + whose row is -1 is skipped), rounded once to `block_out`'s dtype -- bitwise + what the runner's finalize stores. `topk_rows` is the runner's + expanded->permuted index map `[T, k]` int32, `topk_weights` its `[T, k]` + routing weights. The fold needs the combine (`shared`, `gate_logits`). + + Fusion boundary: the routed experts' finalize (optional), the + shared-expert combine (optional), the stream-broadcast + gated accumulate (the old `qwen4_exp_post_inject`) and the grouped gemma + RMS norm of the next site (the old `_grouped_norm`). The routed and shared + experts, the gate projection and the mixer that consumes `normed` are not + inside. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: rows, inject (True), moe (False: shared + gate_logits too), gate + # "column" [T, 1] | "flat" [T], hc_count (4), hidden (2560), dtype (bf16), + # layout "contiguous" | "row_slice" (hyper_input = buf[::2]). + CELLS: tuple[Cell, ...] = ( + tuple( + Cell(why=f"checkpoint 4x2560 bf16, inject, {why}", spec=dict(rows=t)) + for t, why in ( + (1, "decode 1 token"), + (8, "decode batch 8"), + (33, "33 tokens, odd count"), + (128, "decode batch 128 (max concurrency)"), + (2048, "prefill chunk 2048"), + (8192, "prefill chunk 8192 (max_num_tokens)"), + ) + ) + + tuple( + Cell(why=f"norm only (no inject), {why}", spec=dict(rows=t, inject=False)) + for t, why in ((1, "decode 1 token"), (128, "decode 128"), (8192, "prefill 8192")) + ) + + tuple( + Cell(why=f"MoE combine, [T, 1] gate, {why}", spec=dict(rows=t, moe=True)) + for t, why in ( + (1, "decode 1 token"), + (8, "decode batch 8"), + (33, "33 tokens, odd count"), + (128, "decode batch 128 (max concurrency)"), + (2048, "prefill chunk 2048"), + (8192, "prefill chunk 8192"), + (16384, "prefill chunk 16384 (the ladder's max_num_tokens)"), + ) + ) + + tuple( + Cell( + why=f"MoE combine + finalize fold (top-10), {why}", + spec=dict(rows=t, moe=True, fold=10), + ) + for t, why in ( + (2, "decode batch 2"), + (33, "33 tokens, odd count"), + (128, "decode batch 128 (max concurrency)"), + ) + ) + + ( + Cell( + why="MoE combine + finalize fold, three slots mapped to -1, 8 rows", + spec=dict(rows=8, moe=True, fold=10, drop=True), + ), + Cell( + why="MoE combine, flat [T] gate, 17 rows", spec=dict(rows=17, moe=True, gate="flat") + ), + Cell( + why="MoE combine, float16 operands, 64 rows", + spec=dict(rows=64, moe=True, dtype=torch.float16), + ), + Cell(why="MoE combine, zero rows", spec=dict(rows=0, moe=True)), + ) + + ( + Cell( + why="hyper_input a row slice buf[::2], inject, 16 rows", + spec=dict(rows=16, layout="row_slice"), + ), + Cell(why="hc_count 2 at hidden 2560, 64 rows", spec=dict(rows=64, hc_count=2)), + Cell(why="hc_count 8 at hidden 2560, 64 rows", spec=dict(rows=64, hc_count=8)), + Cell(why="hidden 128 (not the checkpoint), 48 rows", spec=dict(rows=48, hidden=128)), + Cell(why="float16 operands, 512 rows", spec=dict(rows=512, dtype=torch.float16)), + Cell(why="zero rows", spec=dict(rows=0)), + ) + ) + + note = """ + Metadata consumed: none; every (row, stream) independent. Call sites: in + the trunk forward of both targets, once after each block (attention / GDN + and MoE) of each layer, normalizing for the next site -- the layer's MoE + mixer, the next layer's attention mixer, or the stack-final mixer -- plus + norm-only calls at the stack's first site and after the PLE add. 97 calls + per forward on the 48-layer trunk; the 48 after a MoE block pass `shared` + and `gate_logits` (in `DecodeTarget` all but the one before the PLE add, + which keeps the unfused combine, and those also `topk_weights` / + `topk_rows`: the runner's finalize folded in). The MTP layer and the PLE + block keep the unfused ops. + + Preconditions: 2-D operands, last-dim stride 1 (row strides free); the + width divisible by `hc_count`; `inject` `[T, hc_count]`, `block_out` and + `shared` `[T, hidden]` of one dtype (with the fold `block_out` is `[P, + hidden]`, any `P`, read through int32 `topk_rows` `[T, k]`); `gate_logits` `T` elements along its + first axis (`[T, 1]` or `[T]`, any row stride); `norm_weight` contiguous + `[width]`. All on one device. `shared` and `gate_logits` come together and + only with `block_out`. + + Numerics: the combine is fp32 with one rounding to `block_out`'s dtype + (bitwise `fused_sigmoid_gate_mul_add`: the same `tl.sigmoid` and fma); the + accumulate is fp32 with one rounding (bitwise what `qwen4_exp_post_inject` + stores for 16-bit operands up to fma contraction); the norm is fp32 with + `rsqrt`, one rounding. Against the float64 reference the combine can land + one ulp of the block output away (its rounding midpoint), which the + accumulate carries scaled by `inject`; the test's band states that. + Graph-capturable: no host sync, outputs allocated per call. + + Launch: one program per (row, stream), stream-minor in a 1-D grid, so a + row's `block_out` / `shared` are read from HBM once and from L2 by the + other `hc_count - 1` programs; with the fold, one program per two streams + of a row (the gathered rows are read and summed half as often). Every program releases programmatic + dependents at its start (`griddepcontrol.launch_dependents`; off with + TRTLLM_ENABLE_PDL=0): a consumer launched with PDL (the decode mixer's + `qwen4_exp_hc_decode_tc_first`) streams its weights meanwhile and waits for + this grid before reading `normed`. A no-op for any other consumer. At a + hidden of 2048 or more a program runs eight warps below + `_FEW_WARPS_MIN_ROWS` rows and two from it; the RMS sum's fp32 reduction + order follows the warp count. + + The decode MoE sites -- the fold below `_FEW_WARPS_MIN_ROWS` rows, bf16 + operands, `hc_count` 4, a hidden in (2048, 4096], 16-byte aligned rows with + row strides a multiple of 16 -- run `_FoldNormKernel` (CuTe DSL) instead: + a row's two CTAs one cluster, CTA j forming half the block output (the + finalize and the combine) into both CTAs' shared memory, then each + accumulating and normalizing two streams, in the Triton kernel's layout and + reduction tree -- bitwise its `hc_out` and `normed`. GB300, graph replay, + 16 rotating sites, us: 128 rows 7.00 -> 5.97, 64 rows 6.63 -> 5.25, 8 rows + 3.59 -> 3.09; it releases dependents at its start as the Triton kernel does. + + Hard limits: row offsets are int64 (`program_id(0)` widened); the grid is + `rows * hc_count` programs along x (2^31 - 1); the column offset + `stream * hidden + i` is int32, bounded by the width. + """ + + def __call__( + self, + hyper_input: torch.Tensor, + block_out: Optional[torch.Tensor], + inject: Optional[torch.Tensor], + norm_weight: torch.Tensor, + eps: float, + hc_count: int, + shared: Optional[torch.Tensor] = None, + gate_logits: Optional[torch.Tensor] = None, + topk_weights: Optional[torch.Tensor] = None, + topk_rows: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + rows, width = hyper_input.shape + hidden = width // hc_count + normed = torch.empty((rows, width), dtype=hyper_input.dtype, device=hyper_input.device) + has_inject = block_out is not None + has_shared = shared is not None + top_k = topk_rows.shape[1] if topk_rows is not None else 0 + hc_out = ( + torch.empty((rows, width), dtype=hyper_input.dtype, device=hyper_input.device) + if has_inject + else hyper_input + ) + if ( + rows > 0 + and top_k + and _fold_cute_ok( + hyper_input, + block_out, + inject, + norm_weight, + hc_count, + shared, + gate_logits, + topk_weights, + topk_rows, + ) + ): + _fold_cute( + hyper_input, + block_out, + inject, + norm_weight, + eps, + hc_count, + shared, + gate_logits, + topk_weights, + topk_rows, + hc_out, + normed, + ) + elif rows > 0: + block = triton.next_power_of_2(hidden) + # The fold gathers TOPK rows to form a row's block output: two streams a + # program form it half as often (when hc_count is even). + spp = 2 if top_k and hc_count % 2 == 0 else 1 + _hc_inject_norm_kernel[(rows * hc_count // spp,)]( + hyper_input, + block_out if has_inject else hyper_input, + inject if has_inject else hyper_input, + shared if has_shared else hyper_input, + gate_logits if has_shared else hyper_input, + topk_weights if top_k else hyper_input, + topk_rows if top_k else hyper_input, + norm_weight, + hc_out, + normed, + hyper_input.stride(0), + block_out.stride(0) if has_inject else 0, + inject.stride(0) if has_inject else 0, + shared.stride(0) if has_shared else 0, + gate_logits.stride(0) if has_shared else 0, + topk_weights.stride(0) if top_k else 0, + topk_rows.stride(0) if top_k else 0, + hc_out.stride(0), + normed.stride(0), + eps, + HC_COUNT=hc_count, + HIDDEN=hidden, + BLOCK=block, + HAS_INJECT=has_inject, + HAS_SHARED=has_shared, + TOPK=top_k, + STREAMS_PER_PROGRAM=spp, + PDL_TRIGGER=_PDL_TRIGGER, + num_warps=(2 if rows >= _FEW_WARPS_MIN_ROWS else 8) if block >= 2048 else 4, + ) + return hc_out, normed + + def reference( + self, + hyper_input: torch.Tensor, + block_out: Optional[torch.Tensor], + inject: Optional[torch.Tensor], + norm_weight: torch.Tensor, + eps: float, + hc_count: int, + shared: Optional[torch.Tensor] = None, + gate_logits: Optional[torch.Tensor] = None, + topk_weights: Optional[torch.Tensor] = None, + topk_rows: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """float64 throughout. The finalize (with `topk_rows`) and the combine + (with `shared`) are each rounded once to `block_out`'s dtype, as the op + states; `hc_out` is the float64 + accumulate rounded once; `normed` is the float64 norm of that rounded + residual, rounded once -- the same dataflow the op states, with no + other intermediate rounding.""" + dtype = hyper_input.dtype + rows, width = hyper_input.shape + hidden = width // hc_count + if block_out is None: + hc_out = hyper_input + else: + if topk_rows is not None: + valid = (topk_rows >= 0).unsqueeze(-1) + picked = block_out.double()[topk_rows.clamp(min=0).long()] + w = topk_weights.double().unsqueeze(-1) + block_out = torch.where(valid, picked * w, 0.0).sum(1).to(block_out.dtype) + if shared is not None: + s = torch.sigmoid(gate_logits.double().reshape(rows, 1)) + block_out = (block_out.double() + s * shared.double()).to(block_out.dtype) + streams = hyper_input.double().unflatten(-1, (hc_count, hidden)) + gated = block_out.double().unsqueeze(-2) * inject.double().unsqueeze(-1) + hc_out = (streams + gated).flatten(-2).to(dtype) + x = hc_out.double().unflatten(-1, (hc_count, hidden)) + x = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) + normed = (x.flatten(-2) * (1.0 + norm_weight.double())).to(dtype) + return hc_out, normed + + def is_valid( + self, + hyper_input: torch.Tensor, + block_out: Optional[torch.Tensor], + inject: Optional[torch.Tensor], + norm_weight: torch.Tensor, + eps: float, + hc_count: int, + shared: Optional[torch.Tensor] = None, + gate_logits: Optional[torch.Tensor] = None, + topk_weights: Optional[torch.Tensor] = None, + topk_rows: Optional[torch.Tensor] = None, + ) -> None: + # The kernel indexes with raw row strides and unit column stride, and + # reads `inject` / `block_out` / `shared` / the gate per row: a broadcast + # (stride-0) or a wrong-width operand would be read out of bounds or + # silently reused. + assert hyper_input.dim() == 2 and hyper_input.stride(1) == 1, "hyper_input [T, width]" + assert (block_out is None) == (inject is None), "block_out and inject come together" + assert (shared is None) == (gate_logits is None), "shared and gate_logits come together" + assert shared is None or block_out is not None, "the combine needs block_out" + assert (topk_weights is None) == (topk_rows is None), ( + "topk_weights and topk_rows come together" + ) + assert topk_rows is None or shared is not None, "the finalize fold needs the combine" + rows, width = hyper_input.shape + assert width % hc_count == 0, f"width {width} not divisible by hc_count {hc_count}" + assert norm_weight.shape == (width,) and norm_weight.is_contiguous(), "norm_weight [width]" + if block_out is not None: + if topk_rows is None: + assert block_out.shape == (rows, width // hc_count) and block_out.stride(1) == 1 + else: + # The permuted FC2 rows: any row count, read through topk_rows. + assert block_out.dim() == 2 and block_out.shape[1] == width // hc_count + assert block_out.stride(1) == 1, "permuted rows [P, hidden]" + assert topk_rows.dtype == torch.int32, "topk_rows int32 (the runner's index map)" + assert topk_rows.dim() == 2 and topk_rows.shape[0] == rows, "topk_rows [T, k]" + assert topk_rows.stride(1) == 1 and topk_weights.stride(1) == 1 + assert topk_weights.shape == topk_rows.shape, "topk_weights [T, k]" + assert inject.shape == (rows, hc_count) and inject.stride(1) == 1 + if shared is not None: + assert shared.shape == (rows, width // hc_count), "shared [T, hidden]" + assert shared.stride(1) == 1, "shared [T, hidden]" + assert shared.dtype == block_out.dtype, "shared in block_out's dtype" + assert gate_logits.shape in ((rows,), (rows, 1)), "gate_logits [T, 1] or [T]" + assert rows <= 1 or gate_logits.stride(0) != 0, "one gate logit per row" + + +qwen4_exp_hc_inject_norm = Qwen4ExpHcInjectNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_lowrank_act.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_lowrank_act.py new file mode 100644 index 000000000000..ba61c79dad97 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_lowrank_act.py @@ -0,0 +1,153 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Hyper-connection mixer activations between its two projections: the SiLU of +the rank-`R` bottleneck and the block-injection weights, from the first +projection's fp32 logits, in one Triton row kernel. + +Tier 5 (Triton written for this target). Used where the up projection runs on +cuBLAS (prefill-sized row counts), between `cublas_mm` (first projection) and +`cublas_mm` (up projection) + `qwen4_exp_hc_gate_mean`. +""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +@triton.jit +def _hc_lowrank_act_kernel( + logits_ptr, + low_ptr, + inject_ptr, + stride_logits, + stride_low, + stride_inject, + RANK: tl.constexpr, + HC: tl.constexpr, + WITH_INJECT: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + k = tl.arange(0, BLOCK) + z = tl.load(logits_ptr + row * stride_logits + k, mask=k < RANK, other=0.0).to(tl.float32) + z = z * (1.0 / HC) + tl.store( + low_ptr + row * stride_low + k, + (z * tl.sigmoid(z)).to(low_ptr.dtype.element_ty), + mask=k < RANK, + ) + if WITH_INJECT: + i = tl.arange(0, HC) + zi = tl.load(logits_ptr + row * stride_logits + RANK + i).to(tl.float32) + tl.store( + inject_ptr + row * stride_inject + i, + (2.0 * tl.sigmoid(zi * (1.0 / HC))).to(inject_ptr.dtype.element_ty), + ) + + +class Qwen4ExpHcLowrankAct(OpWrapper): + """With `R = rank`, `C = hc_count` and `logits` the fp32 `[T, N]` first + projection `normed @ [down; block_inject; 0]^T` (columns `[0, R)` down, + `[R, R + C)` injection when `with_inject`, the rest padding):: + + low = silu(logits[:, :R] / C) # [T, R], `dtype` + inject = 2 * sigmoid(logits[:, R : R + C] / C) # [T, C], `dtype`, only with_inject + + fp32 inside, one rounding per output. Returns fresh `(low, inject)`; + `inject` is `None` exactly when `with_inject` is False. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: rows, inject (True), hc_count (4), rank (320), logits_width (384), dtype (bf16). + CELLS: tuple[Cell, ...] = tuple( + Cell(why=f"checkpoint rank 320 + 4, padded to 384, {why}", spec=dict(rows=t)) + for t, why in ( + (1, "1 row"), + (128, "128 rows"), + (2048, "prefill 2048"), + (8192, "prefill 8192"), + ) + ) + ( + Cell(why="stack-final mixer: no inject, 8192 rows", spec=dict(rows=8192, inject=False)), + Cell(why="unpadded logits [T, R + C], 33 rows", spec=dict(rows=33, logits_width=324)), + Cell(why="hc_count 8, 64 rows", spec=dict(rows=64, hc_count=8)), + Cell(why="float16 outputs, 64 rows", spec=dict(rows=64, dtype=torch.float16)), + Cell(why="zero rows", spec=dict(rows=0)), + ) + + note = """ + Metadata consumed: none. Call sites: the trunk's hyper-connection sites in + `PrefillTarget` (all 97), between the first projection and the up + projection. Preconditions: `logits` fp32, unit column stride, at least + `rank (+ hc_count)` columns; `hc_count` a power of two. Graph-capturable. + """ + + def __call__( + self, + logits: torch.Tensor, + rank: int, + hc_count: int, + with_inject: bool, + dtype: torch.dtype, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + rows = logits.shape[0] + low = torch.empty((rows, rank), dtype=dtype, device=logits.device) + inject = ( + torch.empty((rows, hc_count), dtype=dtype, device=logits.device) + if with_inject + else None + ) + if rows > 0: + _hc_lowrank_act_kernel[(rows,)]( + logits, + low, + inject if with_inject else low, + logits.stride(0), + low.stride(0), + inject.stride(0) if with_inject else 0, + RANK=rank, + HC=hc_count, + WITH_INJECT=with_inject, + BLOCK=triton.next_power_of_2(rank), + num_warps=4, + ) + return low, inject + + def reference( + self, + logits: torch.Tensor, + rank: int, + hc_count: int, + with_inject: bool, + dtype: torch.dtype, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """float64, one rounding per output.""" + lg = logits.double() + low = F.silu(lg[:, :rank] / hc_count).to(dtype) + inject = None + if with_inject: + inject = (2.0 * torch.sigmoid(lg[:, rank : rank + hc_count] / hc_count)).to(dtype) + return low, inject + + def is_valid( + self, + logits: torch.Tensor, + rank: int, + hc_count: int, + with_inject: bool, + dtype: torch.dtype, + ) -> None: + assert logits.dtype == torch.float32 and logits.stride(1) == 1, "fp32 logits, unit stride" + assert logits.shape[1] >= rank + (hc_count if with_inject else 0), "logits too narrow" + assert hc_count & (hc_count - 1) == 0, "hc_count must be a power of two" + + +qwen4_exp_hc_lowrank_act = Qwen4ExpHcLowrankAct() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean.py new file mode 100644 index 000000000000..1ee0d815777f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean.py @@ -0,0 +1,908 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Hyper-connection mixer second half as one tcgen05 GEMM with a custom +epilogue: the up projection `low @ up^T` accumulates in tensor memory and the +epilogue applies the sigmoid gate to the normed stream and averages the +streams, so the `[T, C*H]` gate logits never reach global memory. + +Tier 6 (CuTe DSL written for this target). A persistent warp-specialized +kernel: one TMA warp loads `low` and the up weight, one MMA warp issues +`tcgen05.mma` into a double-buffered TMEM accumulator, four epilogue warps +read the accumulator, load the matching normed elements straight from global +memory and write the mean. An MMA tile is 128 rows by `C * BH` up-projection +outputs: the `BH`-wide hidden slice `[h0, h0 + BH)` of every stream, read +through a strided view of the unmodified `[C*H, R]` weight (no copy, no +re-layout), so the stream reduction stays inside one thread. +""" + +from __future__ import annotations + +import os + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import torch +import torch.nn.functional as F +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.cute_dsl_kernels.blackwell.custom_pipeline import ( + PipelineTmaUmma, + PipelineUmmaAsync, +) +from tensorrt_llm._torch.cute_dsl_kernels.blackwell.utils import make_ptr + +#: Programmatic dependent launch (sm_90+), off with TRTLLM_ENABLE_PDL=0. +_PDL = os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" + + +@dsl_user_op +def _sigmoid(x, *, loc=None, ip=None): + """sigmoid(x) = 0.5 * tanh(0.5 * x) + 0.5: one MUFU op (tanh.approx.f32, + max relative error 2^-10.99) instead of ex2 + rcp.""" + return cutlass.Float32( + llvm.inline_asm( + T.f32(), + [cutlass.Float32(x).ir_value(loc=loc, ip=ip)], + "{ .reg .f32 t; mul.f32 t, $1, 0f3F000000; tanh.approx.f32 t, t; " + "fma.rn.f32 $0, t, 0f3F000000, 0f3F000000; }", + "=f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +# The NVFP4 epilogue (`fp4=True`) is the `torch.ops.trtllm.fp4_quantize` recipe +# (`cvt_warp_fp16_to_fp4`, e4m3 scales) instruction for instruction, every +# multiply an explicit `mul.rn.f32` so that no compiler contraction or +# reassociation can move a rounding. + + +def _asm(ret, args, text, constraints, loc, ip, side_effects=False): + return llvm.inline_asm( + ret, + [a.ir_value(loc=loc, ip=ip) for a in args], + text, + constraints, + has_side_effects=side_effects, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _rcp_approx_ftz(x, *, loc=None, ip=None): + """`rcp.approx.ftz.f32`, the quantizer's `reciprocal_approximate_ftz`.""" + return cutlass.Float32( + _asm(T.f32(), [cutlass.Float32(x)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", loc, ip) + ) + + +@dsl_user_op +def _absmax16(v, *, loc=None, ip=None): + """max |v_i| over 16 fp32 values (exact; NaN-quiet like the bf16 `__hmax2`).""" + body = "abs.f32 m, $1; " + "".join(f"abs.f32 t, ${i}; max.f32 m, m, t; " for i in range(2, 17)) + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Float32(x) for x in v], + "{ .reg .f32 m, t; " + body + "mov.f32 $0, m; }", + "=f" + ",f" * 16, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e4m3_scale(amax, gs, rcp6, *, loc=None, ip=None): + """The block's scale byte: e4m3(gs * (amax * rcp(6))), `cvt.rn.satfinite`.""" + return cutlass.Int32( + _asm( + T.i32(), + [cutlass.Float32(amax), cutlass.Float32(gs), cutlass.Float32(rcp6)], + "{ .reg .f32 t; .reg .b16 h; mul.rn.f32 t, $1, $3; mul.rn.f32 t, $2, t; " + "cvt.rn.satfinite.e4m3x2.f32 h, 0f00000000, t; cvt.u32.u16 $0, h; }", + "=r,f,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _e2m1_out_scale(sf, rcp_gs, amax, *, loc=None, ip=None): + """amax != 0 ? rcp(float(e4m3 sf) * rcp(gs)) : 0, the data multiplier.""" + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Int32(sf), cutlass.Float32(rcp_gs), cutlass.Float32(amax)], + "{ .reg .b16 h, lo, hi; .reg .b32 f2; .reg .f32 s, t; .reg .pred p; " + "cvt.u16.u32 h, $1; cvt.rn.f16x2.e4m3x2 f2, h; mov.b32 {lo, hi}, f2; " + "cvt.f32.f16 s, lo; mul.rn.f32 t, s, $2; rcp.approx.ftz.f32 t, t; " + "setp.neu.f32 p, $3, 0f00000000; selp.f32 $0, t, 0f00000000, p; }", + "=f,r,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _e2m1x8(v, scale, *, loc=None, ip=None): + """Eight values times `scale`, each to e2m1 (`cvt.rn.satfinite`), packed + into 4 bytes; element `2k` in byte k's low nibble.""" + muls = "".join(f"mul.rn.f32 a{i}, ${i + 1}, $9; " for i in range(8)) + cvts = "".join(f"cvt.rn.satfinite.e2m1x2.f32 b{k}, a{2 * k + 1}, a{2 * k}; " for k in range(4)) + return cutlass.Int32( + _asm( + T.i32(), + [cutlass.Float32(x) for x in v] + [cutlass.Float32(scale)], + "{ .reg .f32 a<8>; .reg .b8 b<4>; " + muls + cvts + "mov.b32 $0, {b0, b1, b2, b3}; }", + "=r" + ",f" * 9, + loc, + ip, + ) + ) + + +@dsl_user_op +def _st_global_v4(addr, a, b, c, d, *, loc=None, ip=None): + """One 16-byte `st.global.v4.b32` (a row's packed e2m1 for 32 columns): a + row-per-thread epilogue pays per store instruction, not per byte.""" + _asm( + T.i32(), + [cutlass.Int64(addr)] + [cutlass.Int32(x) for x in (a, b, c, d)], + "{ st.global.v4.b32 [$1], {$2, $3, $4, $5}; mov.u32 $0, 0; }", + "=r,l,r,r,r,r", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _st_global_b16(addr, x, *, loc=None, ip=None): + """One 2-byte store of the low half of `x` (two e4m3 scale bytes).""" + _asm( + T.i32(), + [cutlass.Int64(addr), cutlass.Int32(x)], + "{ .reg .b16 t; cvt.u16.u32 t, $2; st.global.b16 [$1], t; mov.u32 $0, 0; }", + "=r,l,r", + loc, + ip, + side_effects=True, + ) + + +class HcUpGateMeanKernel: + """`mixed[t, h] = 1/C * sum_s sigmoid((low @ up^T)[t, s*H + h]) * normed[t, s*H + h]`. + + Compile-time: `block_h` (hidden columns per tile, 32 or 64), `hc` (C), + `hidden` (H), `rank` (R), `fp4` (also write the NVFP4 of `mixed`: packed + e2m1 `[T, H/2]` and linear e4m3 scales `[T, H/16]` under the global scale + `gs`). Runtime: the row count and the pointers; without `fp4` the `xq`, + `sf` and `gs` pointers are never read. + """ + + def __init__( + self, block_h: int, hc: int, hidden: int, rank: int, num_n_stage: int = 2, fp4: bool = False + ): + assert block_h in (32, 64) and hidden % block_h == 0 and rank % 64 == 0 + self.bh, self.hc, self.hidden, self.rank = block_h, hc, hidden, rank + self.fp4 = fp4 + self.acc_dtype = cutlass.Float32 + self.io_dtype = cutlass.BFloat16 + self.mma_tiler = (128, hc * block_h, 1) + self.epi_cols = 32 + self.cta_group = tcgen05.CtaGroup.ONE + self.cluster_shape_mn = (1, 1) + self.epilogue_warp_id = (0, 1, 2, 3) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.nrm_warp_id = 6 + self.threads_per_cta = 32 * 7 + self.num_n_stage = num_n_stage + self.epilog_sync_bar_id = 1 + self.tmem_alloc_sync_bar_id = 2 + self.num_acc_stage = 2 + + def _setup(self, a_major, b_major): + tiled_mma = utils.sm100.make_trivial_tiled_mma( + self.io_dtype, a_major, b_major, self.acc_dtype, self.cta_group, self.mma_tiler[:2] + ) + inst_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + self.mma_tiler = (self.mma_tiler[0], self.mma_tiler[1], inst_k * 4) + self.cta_tile_shape_mnk = self.mma_tiler + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), (tiled_mma.thr_id.shape,) + ) + self.epi_tile = (self.mma_tiler[0], self.epi_cols) + a_one = utils.sm100.make_smem_layout_a(tiled_mma, self.mma_tiler, self.io_dtype, 1) + b_one = utils.sm100.make_smem_layout_b(tiled_mma, self.mma_tiler, self.io_dtype, 1) + ab_bytes = cute.size_in_bytes(self.io_dtype, a_one) + cute.size_in_bytes( + self.io_dtype, b_one + ) + smem_capacity = utils.get_smem_capacity_in_bytes() + # One normed stage is the 32-column slice of every stream for one tile half. + self.n_box = (self.mma_tiler[0], self.epi_cols) + self.num_n_bytes = self.hc * self.n_box[0] * self.n_box[1] * 2 + self.num_ab_stage = min( + 8, (smem_capacity - 2048 - self.num_n_stage * self.num_n_bytes) // ab_bytes + ) + self.n_pre = min(self.num_ab_stage, self.rank // self.mma_tiler[2]) + self.n_smem_layout_staged = utils.sm100.make_smem_layout_epi( + self.io_dtype, utils.LayoutEnum.ROW_MAJOR, self.n_box, self.num_n_stage * self.hc + ) + self.a_smem_layout_staged = utils.sm100.make_smem_layout_a( + tiled_mma, self.mma_tiler, self.io_dtype, self.num_ab_stage + ) + self.b_smem_layout_staged = utils.sm100.make_smem_layout_b( + tiled_mma, self.mma_tiler, self.io_dtype, self.num_ab_stage + ) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + acc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) + self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(acc_fake, arch="sm_100") + return tiled_mma + + @cute.jit + def __call__( + self, + m: cutlass.Int32, + low_ptr: cute.Pointer, + up_ptr: cute.Pointer, + normed_ptr: cute.Pointer, + mixed: cute.Tensor, + xq_ptr: cute.Pointer, + sf_ptr: cute.Pointer, + gs_ptr: cute.Pointer, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + ): + bh, hc, hid, rank = self.bh, self.hc, self.hidden, self.rank + nt = hid // bh + width = hc * hid + # low [m, R] (K-major A); the up weight [C*H, R] seen as N = (BH, C, H/BH) so + # that one N tile holds hidden slice [h0, h0+BH) of every stream. + a = cute.make_tensor(low_ptr, cute.make_layout((m, rank, 1), stride=(rank, 1, 0))) + b = cute.make_tensor( + up_ptr, + cute.make_layout(((bh, hc, nt), rank, 1), stride=((rank, hid * rank, bh * rank), 1, 0)), + ) + # normed [m, C*H] is staged through shared memory by TMA in [128, 32] boxes; + # mixed [m, H] is written in the accumulator's N order, its stream mode + # with stride 0 (every stream maps onto the one output column). + nrm = cute.make_tensor(normed_ptr, cute.make_layout((m, width), stride=(width, 1))) + mix = cute.make_tensor( + mixed.iterator, cute.make_layout((m, (bh, hc, nt), 1), stride=(hid, (1, 0, bh), 0)) + ) + # NVFP4 outputs (`fp4` only): 8 e2m1 codes per Int32, one e4m3 byte per 16. + xq = cute.make_tensor(xq_ptr, cute.make_layout((m, hid // 8), stride=(hid // 8, 1))) + xsf = cute.make_tensor(sf_ptr, cute.make_layout((m, hid // 16), stride=(hid // 16, 1))) + gsc = cute.make_tensor(gs_ptr, cute.make_layout(1)) + k_major = utils.LayoutEnum.ROW_MAJOR.mma_major_mode() + tiled_mma = self._setup(k_major, k_major) + + a_op = utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id) + a_smem = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_a, tma_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, a, a_smem, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape + ) + b_op = utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id) + b_smem = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, b, b_smem, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape + ) + tma_atom_n, tma_n = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + nrm, + cute.select(self.n_smem_layout_staged, mode=[0, 1]), + self.n_box, + ) + self.num_tma_load_bytes = cute.size_in_bytes(self.io_dtype, a_smem) + cute.size_in_bytes( + self.io_dtype, b_smem + ) + # Rastered along N: consecutive CTAs take consecutive hidden slices of the + # same rows, so the normed reads in flight cover whole rows rather than + # one 64-byte column block of many rows (~2% faster at 8K-16K rows). + sched = utils.PersistentTileSchedulerParams( + ((m + 127) // 128, nt, 1), (*self.cluster_shape_mn, 1), 1, False + ) + grid = utils.StaticPersistentTileScheduler.get_grid_shape(sched, max_active_clusters) + self.kernel( + tiled_mma, + tma_atom_a, + tma_a, + tma_atom_b, + tma_b, + tma_atom_n, + tma_n, + mix, + xq, + xsf, + gsc, + m, + self.cluster_layout_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.n_smem_layout_staged, + self.epi_tile, + sched, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + stream=stream, + use_pdl=_PDL, + ) + + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_n: cute.CopyAtom, + mN_mn: cute.Tensor, + mX_mnl: cute.Tensor, + mQ: cute.Tensor, + mSF: cute.Tensor, + mGS: cute.Tensor, + m: cutlass.Int32, + cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + n_smem_layout_staged: cute.ComposedLayout, + epi_tile: cute.Tile, + sched_params: utils.PersistentTileSchedulerParams, + ): + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + if warp_idx == self.tma_warp_id: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + if warp_idx == self.nrm_warp_id: + cpasync.prefetch_descriptor(tma_atom_n) + tidx, _, _ = cute.arch.thread_idx() + + @cute.struct + class SharedStorage: + ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] + n_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_n_stage * 2] + tmem_dealloc_mbar: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + ab_producer, ab_consumer = PipelineTmaUmma.create( + barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + ).make_participants() + acc_pipeline = PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len(self.epilogue_warp_id) + ), + cta_layout_vmnk=cluster_layout_vmnk, + ) + n_producer, n_consumer = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.n_full_mbar_ptr.data_ptr(), + num_stages=self.num_n_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len(self.epilogue_warp_id) + ), + tx_count=self.num_n_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + ).make_participants() + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_alloc_sync_bar_id, + num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), + ) + tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.epilogue_warp_id[0], + is_two_cta=False, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr, + ) + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + + sA = smem.allocate_tensor( + element_type=self.io_dtype, + layout=a_smem_layout_staged.outer, + byte_alignment=128, + swizzle=a_smem_layout_staged.inner, + ) + sB = smem.allocate_tensor( + element_type=self.io_dtype, + layout=b_smem_layout_staged.outer, + byte_alignment=128, + swizzle=b_smem_layout_staged.inner, + ) + sN = smem.allocate_tensor( + element_type=self.io_dtype, + layout=n_smem_layout_staged.outer, + byte_alignment=128, + swizzle=n_smem_layout_staged.inner, + ) + + gA_mkl = cute.local_tile( + mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + tile_mn = cute.slice_(self.mma_tiler, (None, None, 0)) + gX_mnl = cute.local_tile(mX_mnl, tile_mn, (None, None, None)) + cX_mnl = cute.local_tile( + cute.make_identity_tensor(mX_mnl.shape), tile_mn, (None, None, None) + ) + # normed boxes: [128 rows, 32 columns] at (row tile, column block) + gN_box = cute.local_tile(mN_mn, self.n_box, (None, None)) + k_tile_cnt = cute.size(gA_mkl, mode=[3]) + + thr_mma = tiled_mma.get_slice(0) + tCgA = thr_mma.partition_A(gA_mkl) + tCgB = thr_mma.partition_B(gB_nkl) + tCgX = thr_mma.partition_C(gX_mnl) + tCcX = thr_mma.partition_C(cX_mnl) + + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + 0, + cute.make_layout(1), + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + 0, + cute.make_layout(1), + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + tNsN, tNgN = cpasync.tma_partition( + tma_atom_n, + 0, + cute.make_layout(1), + cute.group_modes(sN, 0, 2), + cute.group_modes(gN_box, 0, 2), + ) + tCrA = tiled_mma.make_fragment_A(sA) + tCrB = tiled_mma.make_fragment_B(sB) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) + + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + tile_sched = utils.StaticPersistentTileScheduler.create( + sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + if warp_idx == self.tma_warp_id: + # Programmatic dependent launch: the first tile's up-weight stages (static) + # stream before the grid-dependency wait, `low` -- the previous kernel's + # output -- after it. Every CTA owns at least one tile (grid <= tiles). + crd = work_tile.tile_idx + tAgA_slice = tAgA[(None, crd[0], None, crd[2])] + tBgB_slice = tBgB[(None, crd[1], None, crd[2])] + ab_producer.reset() + pre = [] + for kk in cutlass.range_constexpr(self.n_pre): + handle = ab_producer.acquire_and_advance() + cute.copy( + tma_atom_b, + tBgB_slice[(None, kk)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + ) + pre.append(handle) + cute.arch.griddepcontrol_wait() + for kk in cutlass.range_constexpr(self.n_pre): + handle = pre[kk] + cute.copy( + tma_atom_a, + tAgA_slice[(None, kk)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + ) + for k_tile in cutlass.range(self.n_pre, k_tile_cnt, 1, unroll=1): + handle = ab_producer.acquire_and_advance() + cute.copy( + tma_atom_a, + tAgA_slice[(None, k_tile)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, k_tile)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + while work_tile.is_valid_tile: + crd = work_tile.tile_idx + tAgA_slice = tAgA[(None, crd[0], None, crd[2])] + tBgB_slice = tBgB[(None, crd[1], None, crd[2])] + ab_producer.reset() + peek = ab_producer.try_acquire() + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + handle = ab_producer.acquire_and_advance(peek) + cute.copy( + tma_atom_a, + tAgA_slice[(None, handle.count)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, handle.count)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + ) + peek = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek = ab_producer.try_acquire() + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + ab_producer.tail() + + if warp_idx == self.nrm_warp_id: + cute.arch.griddepcontrol_wait() + halves = self.bh // self.epi_cols + col_blocks = self.hidden // self.epi_cols + while work_tile.is_valid_tile: + crd = work_tile.tile_idx + for half in cutlass.range_constexpr(halves): + handle = n_producer.acquire_and_advance() + for s in cutlass.range_constexpr(self.hc): + cute.copy( + tma_atom_n, + tNgN[(None, crd[0], s * col_blocks + crd[1] * halves + half)], + tNsN[(None, handle.index * self.hc + s)], + tma_bar_ptr=handle.barrier, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + n_producer.tail() + + if warp_idx == self.mma_warp_id: + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage + ) + while work_tile.is_valid_tile: + tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] + ab_consumer.reset() + peek = ab_consumer.try_wait() + acc_pipeline.producer_acquire(acc_producer_state) + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + handle = ab_consumer.wait_and_advance(peek) + num_kblocks = cute.size(tCrA, mode=[2]) + for kb in cutlass.range(num_kblocks, unroll_full=True): + kcrd = (None, None, kb, handle.index) + cute.gemm(tiled_mma, tCtAcc, tCrA[kcrd], tCrB[kcrd], tCtAcc) + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + handle.release() + peek = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek = ab_consumer.try_wait() + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + acc_pipeline.producer_tail(acc_producer_state) + + if warp_idx < self.mma_warp_id: + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + copy_atom_t2r = utils.sm100.get_tmem_load_op( + self.cta_tile_shape_mnk, + utils.LayoutEnum.ROW_MAJOR, + self.io_dtype, + self.acc_dtype, + epi_tile, + False, + ) + tAcc_epi = cute.flat_divide(tCtAcc_base[((None, None), 0, 0, None)], epi_tile) + tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)]) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + tTR_tAcc_base = thr_copy_t2r.partition_S(tAcc_epi) + tTR_sN = thr_copy_t2r.partition_D(sN) + tTR_gX = thr_copy_t2r.partition_D( + cute.flat_divide(tCgX[((None, None), 0, 0, None, None, None)], epi_tile) + ) + tTR_cX = thr_copy_t2r.partition_D( + cute.flat_divide(tCcX[((None, None), 0, 0, None, None, None)], epi_tile) + ) + frag_shape = tTR_gX[(None, None, None, 0, 0, 0, 0, 0)].shape + tTR_rAcc = cute.make_rmem_tensor(frag_shape, self.acc_dtype) + tTR_rX = cute.make_rmem_tensor(frag_shape, self.io_dtype) + tTR_rSum = cute.make_rmem_tensor(frag_shape, self.acc_dtype) + nfrag = cute.size(tTR_rAcc) + halves = self.bh // self.epi_cols + rN = cute.make_rmem_tensor(frag_shape, self.io_dtype) + inv_hc = cutlass.Float32(1.0 / self.hc) + if cutlass.const_expr(self.fp4): + # A thread owns one row's `epi_cols` consecutive columns (the + # 32dp32b TMEM load), i.e. whole 16-element scale blocks. + gs = mGS[0] + rcp_gs = _rcp_approx_ftz(gs) + rcp6 = _rcp_approx_ftz(cutlass.Float32(6.0)) + rV = cute.make_rmem_tensor((nfrag,), self.acc_dtype) + assert nfrag == 32, "the NVFP4 epilogue stores two scale blocks per row" + q_base = mQ.iterator.toint() + sf_base = mSF.iterator.toint() + + while work_tile.is_valid_tile: + crd = work_tile.tile_idx + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_consumer_state.index)] + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + gX_t = tTR_gX[(None, None, None, None, None, *crd)] + gX_t = cute.group_modes(gX_t, 3, cute.rank(gX_t)) + cX_t = tTR_cX[(None, None, None, None, None, *crd)] + cX_t = cute.group_modes(cX_t, 3, cute.rank(cX_t)) + row_ok = cX_t[(0, 0, 0, 0)][0] < m + + acc_pipeline.consumer_wait(acc_consumer_state) + for half in cutlass.range_constexpr(halves): + n_handle = n_consumer.wait_and_advance() + for s in cutlass.range_constexpr(self.hc): + cute.autovec_copy( + tTR_sN[(None, None, None, n_handle.index * self.hc + s)], rN + ) + cute.copy( + tiled_copy_t2r, + tTR_tAcc[(None, None, None, s * halves + half)], + tTR_rAcc, + ) + for i in cutlass.range_constexpr(nfrag): + gate = _sigmoid(tTR_rAcc[i]) + if cutlass.const_expr(s == 0): + tTR_rSum[i] = gate * rN[i].to(self.acc_dtype) + else: + tTR_rSum[i] = tTR_rSum[i] + gate * rN[i].to(self.acc_dtype) + cute.arch.sync_warp() + n_handle.release() + for i in cutlass.range_constexpr(nfrag): + tTR_rX[i] = (tTR_rSum[i] * inv_hc).to(self.io_dtype) + if row_ok: + cute.autovec_copy(tTR_rX, gX_t[(None, None, None, half)]) + if cutlass.const_expr(self.fp4): + row = cX_t[(0, 0, 0, 0)][0] + blk0 = (crd[1] * self.bh + half * self.epi_cols) // 16 + for i in cutlass.range_constexpr(nfrag): + rV[i] = tTR_rX[i].to(self.acc_dtype) + words, sfs = [], [] + for b in cutlass.range_constexpr(2): + v = [rV[16 * b + j] for j in range(16)] + amax = _absmax16(v) + sf = _e4m3_scale(amax, gs, rcp6) + osc = _e2m1_out_scale(sf, rcp_gs, amax) + words += [_e2m1x8(v[:8], osc), _e2m1x8(v[8:], osc)] + sfs.append(sf) + row64 = cutlass.Int64(row) + _st_global_v4( + q_base + row64 * (self.hidden // 2) + cutlass.Int64(blk0) * 8, + *words, + ) + _st_global_b16( + sf_base + row64 * (self.hidden // 16) + cutlass.Int64(blk0), + sfs[0] | (sfs[1] << 8), + ) + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + cute.arch.barrier( + barrier_id=self.epilog_sync_bar_id, + number_of_threads=32 * len(self.epilogue_warp_id), + ) + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + +_COMPILED: dict = {} + + +def _launcher(block_h: int, hc: int, hidden: int, rank: int, ptrs, fp4: bool = False): + """Compile once per static configuration; later calls reuse the TVM-FFI entry, + which launches on the caller's current stream (graph-capturable). `ptrs` is + `(low, up, normed, mixed tensor, xq, sf, gs)`.""" + key = (block_h, hc, hidden, rank, fp4) + fn = _COMPILED.get(key) + if fn is None: + kern = HcUpGateMeanKernel(block_h, hc, hidden, rank, fp4=fp4) + cptrs = [ + make_ptr(cutlass.BFloat16, p, cute.AddressSpace.gmem, assumed_align=16) + for p in ptrs[:3] + ] + out = from_dlpack(ptrs[3], assumed_align=16).mark_layout_dynamic(leading_dim=1) + qptrs = [ + make_ptr(cutlass.Int32, ptrs[4], cute.AddressSpace.gmem, assumed_align=16), + make_ptr(cutlass.Uint8, ptrs[5], cute.AddressSpace.gmem, assumed_align=2), + make_ptr(cutlass.Float32, ptrs[6], cute.AddressSpace.gmem, assumed_align=4), + ] + fn = cute.compile( + kern, + 128, + *cptrs, + out, + *qptrs, + max_active_clusters=utils.HardwareInfo().get_max_active_clusters(1), + stream=cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--opt-level 2 --enable-tvm-ffi", + ) + _COMPILED[key] = fn + return fn + + +def hc_up_gate_mean( + low: torch.Tensor, up: torch.Tensor, normed: torch.Tensor, hc_count: int, block_h: int +) -> torch.Tensor: + """Launch the fused kernel; `up` is the `[R, C*H]` transposed view of the + `[C*H, R]` up weight, as `cublas_mm(low, up)` takes it.""" + rows, rank = low.shape + hidden = normed.shape[1] // hc_count + mixed = torch.empty((rows, hidden), dtype=normed.dtype, device=normed.device) + if rows > 0: + ptrs = (low.data_ptr(), up.data_ptr(), normed.data_ptr()) + # The NVFP4 pointers are never read without `fp4`. + unused = (mixed.data_ptr(),) * 3 + _launcher(block_h, hc_count, hidden, rank, (*ptrs, mixed, *unused))( + rows, *ptrs, mixed, *unused + ) + return mixed + + +def hc_up_gate_mean_fp4( + low: torch.Tensor, + up: torch.Tensor, + normed: torch.Tensor, + hc_count: int, + global_scale: torch.Tensor, + block_h: int, +): + """`hc_up_gate_mean` plus the NVFP4 of `mixed` in the same epilogue. + Returns `(mixed [T, H], xq [T, H/2] uint8, xsf [T*H/16] uint8)`.""" + rows, rank = low.shape + hidden = normed.shape[1] // hc_count + dev = normed.device + mixed = torch.empty((rows, hidden), dtype=normed.dtype, device=dev) + xq = torch.empty((rows, hidden // 2), dtype=torch.uint8, device=dev) + xsf = torch.empty((rows * (hidden // 16),), dtype=torch.uint8, device=dev) + if rows > 0: + ptrs = (low.data_ptr(), up.data_ptr(), normed.data_ptr()) + qptrs = (xq.data_ptr(), xsf.data_ptr(), global_scale.data_ptr()) + _launcher(block_h, hc_count, hidden, rank, (*ptrs, mixed, *qptrs), fp4=True)( + rows, *ptrs, mixed, *qptrs + ) + return mixed, xq, xsf + + +# Hidden columns per tile. 32 (MMA N = 4 * 32 = 128, five A/B stages) measured +# equal to 64 (N = 256, three stages) at 8K-16K rows and faster at 32..2048 rows. +_BLOCK_H = 32 + + +class Qwen4ExpHcUpGateMean(OpWrapper): + """With `C = hc_count`, `H = width // C`, `low` the `[T, R]` SiLU bottleneck + and `up` the `[R, C*H]` transposed view of the `[C*H, R]` up weight:: + + mixed[t, i] = mean over s of sigmoid((low @ up)[t, s*H + i]) * normed[t, s*H + i] + + The GEMM accumulates in fp32 in tensor memory and is never rounded or + written; sigmoid, product and stream mean are fp32 with one rounding to + `normed`'s dtype. Returns a fresh `[T, H]` tensor. Fusion boundary: the up + projection, the sigmoid gate and the stream mean (what `cublas_mm(low, up)` + + `qwen4_exp_hc_gate_mean` compute as two launches). + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: rows, hc_count (4), hidden (2560), rank (320); bf16 only. + CELLS: tuple[Cell, ...] = tuple( + Cell(why=f"checkpoint 4x2560, rank 320, {why}", spec=dict(rows=t)) + for t, why in ( + (1, "decode 1 row"), + (32, "decode 32 rows"), + (128, "decode 128 rows"), + (2048, "prefill 2048"), + (8192, "prefill 8192"), + (16384, "prefill chunk 16384 (max_num_tokens)"), + ) + ) + ( + Cell(why="ragged tail, 16320 + 70 rows (mixed step)", spec=dict(rows=16390)), + Cell(why="ragged tail, 1000 rows", spec=dict(rows=1000)), + Cell(why="hc_count 2, 300 rows", spec=dict(rows=300, hc_count=2)), + Cell(why="hc_count 8 (32-column tiles), 300 rows", spec=dict(rows=300, hc_count=8)), + Cell(why="hidden 128, rank 128, 200 rows", spec=dict(rows=200, hidden=128, rank=128)), + Cell(why="zero rows", spec=dict(rows=0)), + ) + + note = """ + Tier 6 (CuTe DSL written for this target): a persistent tcgen05 GEMM whose + epilogue applies the gate, so the `[T, C*H]` gate logits never reach HBM; + tiers 1-5 have no GEMM with a custom multi-input epilogue (cuBLAS writes + the logits; the Triton `qwen4_exp_hc_gate_mix` was 2.7x slower than + cuBLAS + gate_mean at prefill rows). Metadata consumed: none. Call sites: + the trunk's hyper-connection mixers (`_hc_mixer`). Preconditions (bf16): + `low` `[T, R]` contiguous, `up` the transposed view of a contiguous + `[C*H, R]` weight, `normed` `[T, C*H]` contiguous, 16-byte aligned + pointers, `R % 64 == 0`, `H % 32 == 0`, `C * 32 <= 256`. Compiles once per + `(block_h, C, H, R)` on first call (~1 s); the row count is a runtime + argument. Graph-capturable (launches on the current stream, no host sync). + Launched with programmatic dependent launch (TRTLLM_ENABLE_PDL=0 turns it + off): each CTA's first tile streams its up-weight stages before + `griddepcontrol.wait` and `low` / `normed` after it, so behind a producer + that releases dependents early (`qwen4_exp_hc_decode_tc_first`) the weight + read overlaps the producer. + """ + + def __call__( + self, low: torch.Tensor, up: torch.Tensor, normed: torch.Tensor, hc_count: int + ) -> torch.Tensor: + return hc_up_gate_mean(low, up, normed, hc_count, _BLOCK_H) + + def reference( + self, low: torch.Tensor, up: torch.Tensor, normed: torch.Tensor, hc_count: int + ) -> torch.Tensor: + """float64, one rounding.""" + hidden = normed.shape[1] // hc_count + gate = torch.sigmoid(F.linear(low.double(), up.double().t())).unflatten( + -1, (hc_count, hidden) + ) + x = normed.double().unflatten(-1, (hc_count, hidden)) + return (gate * x).mean(dim=-2).to(normed.dtype) + + def is_valid( + self, low: torch.Tensor, up: torch.Tensor, normed: torch.Tensor, hc_count: int + ) -> None: + rows, rank = low.shape + width = normed.shape[1] + assert low.dtype == up.dtype == normed.dtype == torch.bfloat16, "bf16 operands" + assert normed.shape[0] == rows and up.shape == (rank, width), "shapes" + assert low.stride() == (rank, 1) and normed.stride() == (width, 1), "contiguous rows" + assert up.stride() == (1, rank), "up must be the transposed view of a [C*H, R] weight" + assert width % hc_count == 0 and (width // hc_count) % 32 == 0 and hc_count * 32 <= 256 + assert rank % 64 == 0, "rank must be a multiple of 64" + for t in (low, up, normed): + assert t.data_ptr() % 16 == 0, "16-byte aligned operands" + + +qwen4_exp_hc_up_gate_mean = Qwen4ExpHcUpGateMean() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean_fp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean_fp4.py new file mode 100644 index 000000000000..654d5263d1c3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean_fp4.py @@ -0,0 +1,124 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Hyper-connection mixer second half at an MoE site: `qwen4_exp_hc_up_gate_mean` +(the up projection as a tcgen05 GEMM, sigmoid gate and stream mean in its +epilogue) whose epilogue also writes the routed experts' NVFP4 input. + +Tier 6 (CuTe DSL written for this target): the `fp4=True` build of +`qwen4_exp_hc_up_gate_mean`'s kernel. The routed experts read the mixer's output +only as NVFP4 (packed e2m1 + one e4m3 scale per 16 elements, linear scale +layout); the bf16 mix is still written for the router and the shared expert. +Each epilogue thread already holds one row's 32 consecutive mix values, i.e. +two whole scale blocks, so the quantization needs no data exchange and removes +the `fp4_quantize` launch and its re-read of the bf16 activation. +""" + +from __future__ import annotations + +from typing import Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .qwen4_exp_hc_decode_mix_fp4 import SF_VEC, nvfp4_reference +from .qwen4_exp_hc_up_gate_mean import _BLOCK_H, hc_up_gate_mean_fp4, qwen4_exp_hc_up_gate_mean + + +class Qwen4ExpHcUpGateMeanFp4(OpWrapper): + """With `C = hc_count`, `H = width // C`, `low` the `[T, R]` SiLU bottleneck, + `up` the `[R, C*H]` transposed view of the `[C*H, R]` up weight:: + + mixed[t, i] = mean over s of sigmoid((low @ up)[t, s*H + i]) * normed[t, s*H + i] + xq, xsf = NVFP4(mixed) # e2m1 pairs, e4m3 scale per 16, linear layout + + `mixed` is bitwise `qwen4_exp_hc_up_gate_mean`'s (same kernel, same + epilogue arithmetic); the quantization reads that bf16-rounded value, so + `xq` / `xsf` are bitwise `torch.ops.trtllm.fp4_quantize(mixed, gs, 16, + False, False)`. Returns fresh `(mixed [T, H], xq [T, H/2] uint8, xsf + [T*H/16] uint8)`. `global_scale` is the 1-element fp32 NVFP4 global scale + (`1 / input_scale`). Fusion boundary: the up projection, the sigmoid + gate, the stream mean and the NVFP4 quantization of the mean. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: rows, hc_count (4), hidden (2560), rank (320); bf16 only. + CELLS: tuple[Cell, ...] = tuple( + Cell(why=f"checkpoint 4x2560, rank 320, {why}", spec=dict(rows=t)) + for t, why in ( + (2, "decode 2 rows (smallest fused-up batch)"), + (32, "decode 32 rows"), + (128, "decode 128 rows (max_con)"), + (2048, "prefill 2048"), + (8192, "prefill 8192"), + (16384, "prefill chunk 16384 (max_num_tokens)"), + ) + ) + ( + Cell(why="1 row (below the decode threshold)", spec=dict(rows=1)), + Cell(why="ragged tail, 16320 + 70 rows (mixed step)", spec=dict(rows=16390)), + Cell(why="ragged tail, 1000 rows", spec=dict(rows=1000)), + Cell(why="hc_count 2, 300 rows", spec=dict(rows=300, hc_count=2)), + Cell(why="hc_count 8 (32-column tiles), 300 rows", spec=dict(rows=300, hc_count=8)), + Cell(why="hidden 128, rank 128, 200 rows", spec=dict(rows=200, hidden=128, rank=128)), + Cell(why="zero rows", spec=dict(rows=0)), + ) + + note = """ + Tier 6 (CuTe DSL written for this target): `qwen4_exp_hc_up_gate_mean`'s + kernel compiled with its NVFP4 epilogue; tiers 1-5 have no GEMM with this + multi-input epilogue, and the fused up projection itself is tier 6. + Metadata consumed: none. Call sites: the 48 MoE-site mixers (prefill, and + decode from `_HC_FUSED_UP_MIN_ROWS` rows); `xq`/`xsf` feed + `fp4_block_scale_moe_renorm_routed` (linear scales), `mixed` the router and + the shared expert. Preconditions: those of `qwen4_exp_hc_up_gate_mean` + (bf16, `R % 64 == 0`, `H % 32 == 0`, `C * 32 <= 256`, 16-byte aligned + operands) and `global_scale` one fp32 element on the device. Compiles once + per `(C, H, R)` on first call; graph-capturable (current stream, no host + sync). + + Numerics: the quantization is the trtllm NVFP4 recipe instruction for + instruction (`rcp.approx.ftz.f32` reciprocals, `cvt.rn.satfinite` to e4m3 + and to e2m1, every multiply `mul.rn.f32`); the GPU test checks it bitwise + against `torch.ops.trtllm.fp4_quantize` on the returned `mixed`. + """ + + def __call__( + self, + low: torch.Tensor, + up: torch.Tensor, + normed: torch.Tensor, + hc_count: int, + global_scale: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return hc_up_gate_mean_fp4(low, up, normed, hc_count, global_scale, _BLOCK_H) + + def reference( + self, + low: torch.Tensor, + up: torch.Tensor, + normed: torch.Tensor, + hc_count: int, + global_scale: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """`mixed` in float64 with one rounding, then the NVFP4 of that rounded + value in native torch.""" + mixed = qwen4_exp_hc_up_gate_mean.reference(low, up, normed, hc_count) + return (mixed,) + nvfp4_reference(mixed, global_scale) + + def is_valid( + self, + low: torch.Tensor, + up: torch.Tensor, + normed: torch.Tensor, + hc_count: int, + global_scale: torch.Tensor, + ) -> None: + qwen4_exp_hc_up_gate_mean.is_valid(low, up, normed, hc_count) + assert (normed.shape[1] // hc_count) % SF_VEC == 0, "hidden must be whole 16-blocks" + # The kernel loads element 0 only; a per-row scale would be applied as row 0's. + assert global_scale.numel() == 1 and global_scale.dtype == torch.float32 + assert global_scale.is_cuda + + +qwen4_exp_hc_up_gate_mean_fp4 = Qwen4ExpHcUpGateMeanFp4() diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_first.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_first.py new file mode 100644 index 000000000000..9126b995746e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_first.py @@ -0,0 +1,450 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qwen4_exp_hc_decode_first entry. + +Cells against the float64 reference: `hc_out` within one ulp (as +`qwen4_exp_hc_inject_norm`), `part` within `2^-16` of the reference (relative +to its largest element, plus the same relative term), `ss` to fp32 accuracy. +The MoE-combine cells widen `hc_out`'s band by the combine's own rounding and +check `part` / `ss` against the projection of the `hc_out` the op stored (see +`_compare_moe`). Beyond the cells: wrong formulas land outside, bitwise +agreement with the combine + this op composition, the weight-orientation +guard, determinism 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_hc_decode_first import ( # noqa: E501 + COMBINE_CONFIG, + FOLD_CONFIG, + launch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_decode_first import ( # noqa: E501 + qwen4_exp_hc_decode_first as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp4_block_scale_moe_renorm_permuted import ( # noqa: E501 + finalize, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fused_sigmoid_gate_mul_add import ( # noqa: E501 + fused_sigmoid_gate_mul_add, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "qwen4_exp_hc_decode_first requires a CUDA device" + +DEV = "cuda" +HC_COUNT, HIDDEN, RANK = 4, 2560, 320 +# This checkpoint's layer-0 hyper-connection element statistics. +DOWN_RMS, INJECT_RMS, GAMMA_RMS, GAMMA_ABSMAX = 0.0296, 0.0107, 0.477, 6.625 +PART_EPS = 2.0**-16 + + +def _band(got, want, eps) -> float: + """Worst element's distance as a fraction of `rtol = eps`, `atol = eps * max|want|`.""" + w = want.double() + return ((got.double() - w).abs() / (eps * w.abs().max().item() + eps * w.abs())).max().item() + + +def _ulp(got, want) -> float: + eps, tiny = torch.finfo(want.dtype).eps, torch.finfo(want.dtype).tiny + w = want.double() + return ((got.double() - w).abs() / (eps * w.abs() + tiny)).max().item() + + +def _compare(out, ref) -> None: + """`hc_out` one ulp: an fp32 accumulate rounded once. `part`: a 2560-term + fp32 dot per stream, so ~2^-20 relative to its scale is the expected error; + the gate is 2^-16. `ss`: a 2560-term fp32 sum of squares.""" + (hc, part, ss), (ref_hc, ref_part, ref_ss) = out, ref + assert hc.dtype == ref_hc.dtype and hc.shape == ref_hc.shape + assert part.dtype == torch.float32 and part.shape == ref_part.shape and part.stride(2) == 1 + assert ss.dtype == torch.float32 and ss.shape == ref_ss.shape + if not hc.numel(): + return + assert _ulp(hc, ref_hc) <= 1.0, f"hc_out {_ulp(hc, ref_hc):.4g} ulp" + assert _band(part, ref_part, PART_EPS) <= 1.0, f"part {_band(part, ref_part, PART_EPS):.4g}x" + assert torch.allclose(ss, ref_ss, rtol=1e-5, atol=0), "ss" + + +def _combine(block_out, shared, gate_logits) -> torch.Tensor: + """The MoE block output in float64, rounded once to `block_out`'s dtype.""" + s = torch.sigmoid(gate_logits.double().reshape(-1, 1)) + return (block_out.double() + s * shared.double()).to(block_out.dtype) + + +def _routed(block_out, topk_weights, topk_rows) -> torch.Tensor: + """The routed row: `block_out` itself, or with the fold the float64 + top-k weighted sum of the permuted rows rounded once to their dtype.""" + if topk_rows is None: + return block_out + valid = (topk_rows >= 0).unsqueeze(-1) + rows = block_out.double()[topk_rows.clamp(min=0).long()] + return ( + torch.where(valid, rows * topk_weights.double().unsqueeze(-1), 0.0) + .sum(1) + .to(block_out.dtype) + ) + + +def _hc_band_ratio( + hc, ref_hc, block_out, inject, shared, gate, topk_weights=None, topk_rows=None +) -> float: + """`hc_out` against the reference: one ulp of `hc_out` plus one ulp of the + combine times `|inject|` (the op's fp32 combine and the reference's float64 + one can round to neighbouring values at a midpoint); with the fold, one ulp + of the routed row times `|inject|` more (the op's fp32 FMA finalize and the + reference's exact sum can round apart the same way).""" + feps, tiny = torch.finfo(ref_hc.dtype).eps, torch.finfo(ref_hc.dtype).tiny + routed = _routed(block_out, topk_weights, topk_rows) + b = _combine(routed, shared, gate).double().abs() + if topk_rows is not None: + b = b + routed.double().abs() + band = feps * ref_hc.double().abs() + ( + feps * b.unsqueeze(-2) * inject.double().abs().unsqueeze(-1) + ).flatten(-2) + return ((hc.double() - ref_hc.double()).abs() / (band + tiny)).max().item() + + +def _compare_moe(args, out, ref) -> None: + """MoE-combine cells: `hc_out` in `_hc_band_ratio`'s band of the float64 + result; `part` and `ss` as `_compare` against `reference(hc_out, None, + None, ...)` -- the projection of the residual the op stored.""" + hyper_input, block_out, inject, gamma, first, c, n, shared, gate = args[:9] + topk_weights, topk_rows = args[9:] if len(args) > 9 else (None, None) + (hc, part, ss), (ref_hc, _, _) = out, ref + assert hc.dtype == ref_hc.dtype and hc.shape == ref_hc.shape + if not hc.numel(): + return + ratio = _hc_band_ratio(hc, ref_hc, block_out, inject, shared, gate, topk_weights, topk_rows) + assert ratio <= 1.0, f"hc_out {ratio:.4g}x the band" + _compare((hc, part, ss), op.reference(hc, None, None, gamma, first, c, n)) + + +def _build(spec, seed): + rows = spec["rows"] + hc_count = spec.get("hc_count", HC_COUNT) + hidden = spec.get("hidden", HIDDEN) + dtype = spec.get("dtype", torch.bfloat16) + width = hc_count * hidden + n_cols = spec.get("n_cols", RANK + hc_count) + g = torch.Generator(device=DEV).manual_seed(seed) + + def normal(*shape, scale): + return torch.randn(*shape, generator=g, device=DEV) * scale + + hyper_input = normal(rows, width, scale=3.0).to(dtype) + if spec.get("layout") == "row_slice": + buf = normal(2 * rows, width, scale=1.0).to(dtype) + buf[::2] = hyper_input + hyper_input = buf[::2] + gamma = normal(width, scale=GAMMA_RMS) + gamma[0], gamma[width // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX + gamma = gamma.to(dtype) + n_rows = spec.get("n_rows", -(-n_cols // 128) * 128) + weight = torch.zeros(n_rows, width, device=DEV) + weight[:n_cols] = normal(n_cols, width, scale=DOWN_RMS) + weight[n_cols - hc_count : n_cols] *= INJECT_RMS / DOWN_RMS + first = weight.to(dtype).t() + block_out = inject = shared = gate_logits = None + if spec.get("inject", True): + block_out = normal(rows, hidden, scale=2.0).to(dtype) + inject = (2.0 * torch.sigmoid(normal(rows, hc_count, scale=1.0))).to(dtype) + if spec.get("moe", False): + shared = normal(rows, hidden, scale=2.0).to(dtype) + gate_logits = normal(rows, 1, scale=2.0).to(dtype) + if spec.get("gate") == "flat": + gate_logits = gate_logits.reshape(rows) + top_k = spec.get("fold", 0) + if not top_k: + return hyper_input, block_out, inject, gamma, first, hc_count, n_cols, shared, gate_logits + # The runner's unfinalized output: padded permuted rows (garbage where no + # slot maps), a distinct row per slot, softmax-like bf16 weights. + n_perm = rows * top_k + 37 + block_out = normal(n_perm, hidden, scale=2.0).to(dtype) + order = torch.randperm(n_perm, generator=g, device=DEV)[: rows * top_k] + topk_rows = order.reshape(rows, top_k).to(torch.int32) + if spec.get("drop"): + topk_rows[::2, -3:] = -1 + topk_weights = torch.softmax(normal(rows, top_k, scale=1.0), -1).to(dtype) + return ( + hyper_input, + block_out, + inject, + gamma, + first, + hc_count, + n_cols, + shared, + gate_logits, + topk_weights, + topk_rows, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) + before = [a.clone() if isinstance(a, torch.Tensor) else a for a in args] + out = op(*args) + for a, b in zip(args, before): + assert not isinstance(a, torch.Tensor) or torch.equal(a, b), "every input is read-only" + if args[1] is None: + assert out[0] is args[0], "norm-only returns hyper_input itself" + if args[7] is None: + _compare(out, op.reference(*args)) + else: + _compare_moe(args, out, op.reference(*args)) + + +def test_measured_distance_from_float64() -> None: + for rows in (1, 16): + args = _build(dict(rows=rows), seed=rows) + (hc, part, ss), (rh, rp, rs) = op(*args), op.reference(*args) + print( + f"\n rows {rows}: hc {_ulp(hc, rh):.3g} ulp, part {_band(part, rp, PART_EPS):.4f}x, " + f"ss rel {((ss - rs).abs() / rs).max().item():.3g}" + ) + + +def test_deterministic() -> None: + args = _build(dict(rows=4), seed=3) + for a, b in zip(op(*args), op(*args)): + assert torch.equal(a, b) + + +def test_wrong_variants_land_outside_the_band() -> None: + """Each mis-reading keeps the shapes; each lands > 5x outside. The streams + carry different scales so a stream mix-up is visible.""" + args = list(_build(dict(rows=8), seed=5)) + scales = torch.tensor([0.5, 1.0, 2.0, 4.0], device=DEV).repeat_interleave(HIDDEN) + args[0] = (args[0].float() * scales).to(torch.bfloat16) + hyper_input, block_out, inject, gamma, first, c, n, _, _ = args + want_hc, want_part, _ = op.reference(*args) + pre = hyper_input.double().unflatten(-1, (c, HIDDEN)) + hc = want_hc.double().unflatten(-1, (c, HIDDEN)) + g = (1.0 + gamma.double()).unflatten(-1, (c, HIDDEN)) + w = first[:, :n].double().unflatten(0, (c, HIDDEN)) + + def proj(y): + return torch.einsum("tsi,sin->stn", y, w) + + variants = { + "w instead of 1 + w": proj(hc * (g - 1.0)), + "pre-inject stream": proj(pre * g), + "no (1 + w)": proj(hc), + "streams summed into every partial": proj(hc * g).sum(0, keepdim=True).expand(c, -1, -1), + } + for name, value in variants.items(): + ratio = _band(value.float(), want_part, PART_EPS) + print(f"\n part {name:<36} {ratio:10.4g}x") + assert ratio > 5.0, name + stream_minor = pre.flatten(-2) + ( + block_out.double().unsqueeze(-1) * inject.double().unsqueeze(-2) + ).flatten(-2) + ratio = _ulp(stream_minor.to(torch.bfloat16), want_hc) + print(f" hc_out stream-minor inject {ratio:10.4g} ulp") + assert ratio > 5.0 + + +def test_guard_refuses_a_row_major_weight() -> None: + """A contiguous `[width, N]` weight is read with rows and columns swapped + and returns plausible numbers unguarded.""" + args = list(_build(dict(rows=2), seed=6)) + want = op.reference(*args)[1] + args[4] = args[4].contiguous() + assert _band(op(*args)[1], want, PART_EPS) > 5.0, "a row-major weight should answer wrongly" + with validating(op), pytest.raises(AssertionError): + op(*args) + + +def test_cuda_graph_capture_and_replay() -> None: + args = list(_build(dict(rows=1), seed=8)) + static_in = args[0] + 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): + graph_out = op(*args) + fresh = _build(dict(rows=1), seed=18)[0] + static_in.copy_(fresh) + graph.replay() + torch.cuda.synchronize() + for a, b in zip(graph_out, op(fresh, *args[1:])): + assert torch.equal(a, b) + + +@pytest.mark.parametrize("rows,n_cols", [(1, RANK + HC_COUNT), (1, RANK), (16, RANK + HC_COUNT)]) +def test_moe_combine_matches_the_composition_it_replaces(rows, n_cols) -> None: + """Against `fused_sigmoid_gate_mul_add` (in place on the routed output, as + the forward called it) and then this op without the combine, at the tile + the combine path launches with: bitwise. (The default tile sums `part` in + another fp32 order; the cells bound both against the reference.)""" + args = _build(dict(rows=rows, moe=True, n_cols=n_cols), seed=200 + rows) + hyper_input, block_out, inject, gamma, first, c, n, shared, gate = args + fused = op(*args) + combined = fused_sigmoid_gate_mul_add(block_out.clone(), gate, shared) + hc, part, ss = launch(hyper_input, combined, inject, gamma, first, c, n, config=COMBINE_CONFIG) + for name, a, b in zip(("hc_out", "part", "ss"), fused, (hc, part[:, :, :n], ss)): + assert torch.equal(a, b), name + default = op(hyper_input, combined, inject, gamma, first, c, n) + assert torch.equal(fused[0], default[0]), "hc_out does not depend on the tile" + print( + f"\n rows {rows} cols {n_cols}: vs the default tile, part bitwise " + f"{(fused[1] == default[1]).float().mean().item():.4f} " + f"({_band(fused[1], default[1], PART_EPS):.3g}x the band apart), ss bitwise " + f"{(fused[2] == default[2]).float().mean().item():.4f}" + ) + + +def test_moe_wrong_variants_land_outside_the_band() -> None: + """Each mis-reading of the combine lands > 5x outside `_hc_band_ratio`.""" + args = _build(dict(rows=8, moe=True), seed=11) + hyper_input, block_out, inject, gamma, first, c, n, shared, gate = args + want_hc = op.reference(*args)[0] + r, sh, s = block_out.double(), shared.double(), torch.sigmoid(gate.double()) + variants = { + "no sigmoid": r + gate.double() * sh, + "gate on the routed half": s * r + sh, + "gate dropped": r + sh, + "shared dropped": r, + } + x = hyper_input.double().unflatten(-1, (c, HIDDEN)) + for name, blk in variants.items(): + wrong = (x + blk.unsqueeze(-2) * inject.double().unsqueeze(-1)).flatten(-2) + ratio = _hc_band_ratio(wrong.to(torch.bfloat16), want_hc, block_out, inject, shared, gate) + print(f"\n hc_out {name:<26} {ratio:10.4g}x the band") + assert ratio > 5.0, name + + +def test_guard_refuses_a_partial_combine() -> None: + args = list(_build(dict(rows=2, moe=True), seed=12)) + hyper_input, block_out, inject, gamma, first, c, n, shared, gate = args + with validating(op), pytest.raises(AssertionError): + op(hyper_input, block_out, inject, gamma, first, c, n, shared, None) + with validating(op), pytest.raises(AssertionError): + op(hyper_input, None, None, gamma, first, c, n, shared, gate) + + +def test_moe_cuda_graph_capture_and_replay() -> None: + args = list(_build(dict(rows=1, moe=True), seed=13)) + 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): + graph_out = op(*args) + fresh = _build(dict(rows=1, moe=True), seed=23) + for i in (0, 1, 7, 8): + args[i].copy_(fresh[i]) + graph.replay() + torch.cuda.synchronize() + for a, b in zip(graph_out, op(*args)): + assert torch.equal(a, b) + + +@pytest.mark.parametrize( + "rows,n_cols,drop", + [ + (1, RANK + HC_COUNT, False), + (1, RANK, False), + (4, RANK + HC_COUNT, True), + (16, RANK + HC_COUNT, False), + ], +) +def test_fold_matches_finalize_then_combine(rows, n_cols, drop) -> None: + """The fold against the runner's finalize (emulated exactly: fp32 FMAs in + slot order, one rounding; the MoE entry's test pins that emulation to the + runner bitwise) followed by this op's combine path at the same tile: + `hc_out` and `ss` bitwise; `part` bitwise from 2 rows, and at 1 row (where + the two Triton specializations reduce the dot in a different fp32 order) + within 2^-22 of it relative to its scale. And `hc_out` bitwise the combine + path at its own tile.""" + args = _build(dict(rows=rows, moe=True, fold=10, drop=drop, n_cols=n_cols), seed=300 + rows) + hyper_input, perm, inject, gamma, first, c, n, shared, gate, topk_w, topk_rows = args + fused = op(*args) + routed = finalize(perm, topk_w, topk_rows) + hc, part, ss = launch( + hyper_input, routed, inject, gamma, first, c, n, shared, gate, config=FOLD_CONFIG + ) + part = part[:, :, :n] + assert torch.equal(fused[0], hc), "hc_out" + assert torch.equal(fused[2], ss), "ss" + if rows > 1: + assert torch.equal(fused[1], part), "part" + else: + assert _band(fused[1], part, 2.0**-22) <= 1.0, "part" + print( + f"\n rows {rows} cols {n_cols}: part vs finalize + combine bitwise " + f"{(fused[1] == part).float().mean().item():.4f} ({_band(fused[1], part, PART_EPS):.3g}x the band apart)" + ) + combined = op(hyper_input, routed, inject, gamma, first, c, n, shared, gate) + assert torch.equal(fused[0], combined[0]), "hc_out does not depend on the tile" + + +def test_fold_wrong_variants_land_outside_the_band() -> None: + """Each mis-reading of the finalize lands > 5x outside `_hc_band_ratio`.""" + args = _build(dict(rows=8, moe=True, fold=10), seed=17) + hyper_input, perm, inject, gamma, first, c, n, shared, gate, topk_w, topk_rows = args + want_hc = op.reference(*args)[0] + rows = perm.double()[topk_rows.long()] + w = topk_w.double().unsqueeze(-1) + s = torch.sigmoid(gate.double()).reshape(-1, 1) + variants = { + "unweighted sum": rows.sum(1), + "first slot only": rows[:, 0] * w[:, 0], + "weights by the wrong slot": (rows * w.flip(1)).sum(1), + "token-major rows (no index map)": (perm.double()[: 8 * 10].reshape(8, 10, -1) * w).sum(1), + } + x = hyper_input.double().unflatten(-1, (c, HIDDEN)) + for name, routed in variants.items(): + blk = routed.to(torch.bfloat16).double() + s * shared.double() + wrong = (x + blk.unsqueeze(-2) * inject.double().unsqueeze(-1)).flatten(-2) + ratio = _hc_band_ratio( + wrong.to(torch.bfloat16), want_hc, perm, inject, shared, gate, topk_w, topk_rows + ) + print(f"\n hc_out {name:<34} {ratio:10.4g}x the band") + assert ratio > 5.0, name + + +def test_guard_refuses_a_partial_fold() -> None: + args = list(_build(dict(rows=2, moe=True, fold=10), seed=19)) + hyper_input, perm, inject, gamma, first, c, n, shared, gate, topk_w, topk_rows = args + with validating(op), pytest.raises(AssertionError, match="needs the combine"): + op(hyper_input, perm, inject, gamma, first, c, n, None, None, topk_w, topk_rows) + with validating(op), pytest.raises(AssertionError, match="come together"): + op(hyper_input, perm, inject, gamma, first, c, n, shared, gate, None, topk_rows) + with validating(op), pytest.raises(AssertionError, match="int32"): + op(hyper_input, perm, inject, gamma, first, c, n, shared, gate, topk_w, topk_rows.long()) + + +def test_fold_cuda_graph_capture_and_replay() -> None: + args = list(_build(dict(rows=1, moe=True, fold=10), seed=29)) + 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): + graph_out = op(*args) + fresh = _build(dict(rows=1, moe=True, fold=10), seed=31) + for i in (0, 1, 2, 7, 8, 9, 10): + args[i].copy_(fresh[i]) + graph.replay() + torch.cuda.synchronize() + for a, b in zip(graph_out, op(*args)): + assert torch.equal(a, b) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py new file mode 100644 index 000000000000..6a419915ba65 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py @@ -0,0 +1,241 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qwen4_exp_hc_decode_first_cute entry. + +Cells against the Triton entry's float64 reference, in the Triton entry's bands: +`hc_out` within one ulp (plus the combine's and the finalize's own rounding at the +MoE sites), `part` within `2^-16` of the reference relative to its scale, `ss` to +fp32 accuracy. Beyond the cells: `hc_out` bitwise the Triton entry's (the same fp32 +arithmetic), `part` / `ss` inside the band of it; a stale-read check under CUDA-graph +replay with PDL, where a kernel that read its activations before +`griddepcontrol.wait` would see the previous replay's; determinism; 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.hc.qwen4_exp_hc_decode_first import ( # noqa: E501 + qwen4_exp_hc_decode_first as triton_op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_decode_first_cute import ( # noqa: E501 + qwen4_exp_hc_decode_first_cute as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "qwen4_exp_hc_decode_first_cute requires a CUDA device" + +DEV = "cuda" +HC_COUNT, HIDDEN, RANK = 4, 2560, 320 +DOWN_RMS, INJECT_RMS, GAMMA_RMS, GAMMA_ABSMAX = 0.0296, 0.0107, 0.477, 6.625 +PART_EPS = 2.0**-16 + + +def _band(got, want, eps) -> float: + w = want.double() + return ((got.double() - w).abs() / (eps * w.abs().max().item() + eps * w.abs())).max().item() + + +def _ulp(got, want) -> float: + eps, tiny = torch.finfo(want.dtype).eps, torch.finfo(want.dtype).tiny + w = want.double() + return ((got.double() - w).abs() / (eps * w.abs() + tiny)).max().item() + + +def _compare(out, ref) -> None: + (hc, part, ss), (ref_hc, ref_part, ref_ss) = out, ref + assert hc.dtype == ref_hc.dtype and hc.shape == ref_hc.shape + assert part.dtype == torch.float32 and part.shape == ref_part.shape and part.stride(2) == 1 + assert ss.dtype == torch.float32 and ss.shape == ref_ss.shape + assert _ulp(hc, ref_hc) <= 1.0, f"hc_out {_ulp(hc, ref_hc):.4g} ulp" + assert _band(part, ref_part, PART_EPS) <= 1.0, f"part {_band(part, ref_part, PART_EPS):.4g}x" + assert torch.allclose(ss, ref_ss, rtol=1e-5, atol=0), "ss" + + +def _routed(block_out, topk_weights, topk_rows): + if topk_rows is None: + return block_out + valid = (topk_rows >= 0).unsqueeze(-1) + rows = block_out.double()[topk_rows.clamp(min=0).long()] + return ( + torch.where(valid, rows * topk_weights.double().unsqueeze(-1), 0.0) + .sum(1) + .to(block_out.dtype) + ) + + +def _hc_band_ratio( + hc, ref_hc, block_out, inject, shared, gate, topk_weights=None, topk_rows=None +) -> float: + """The Triton entry's `hc_out` band at the MoE sites: one ulp of `hc_out`, plus one + ulp of the combine (and with the fold, of the routed row) times `|inject|`.""" + feps, tiny = torch.finfo(ref_hc.dtype).eps, torch.finfo(ref_hc.dtype).tiny + routed = _routed(block_out, topk_weights, topk_rows) + s = torch.sigmoid(gate.double().reshape(-1, 1)) + b = (routed.double() + s * shared.double()).to(routed.dtype).double().abs() + if topk_rows is not None: + b = b + routed.double().abs() + band = feps * ref_hc.double().abs() + ( + feps * b.unsqueeze(-2) * inject.double().abs().unsqueeze(-1) + ).flatten(-2) + return ((hc.double() - ref_hc.double()).abs() / (band + tiny)).max().item() + + +def _compare_moe(args, out, ref) -> None: + hyper_input, block_out, inject, gamma, first, c, n, shared, gate = args[:9] + topk_weights, topk_rows = args[9:] if len(args) > 9 else (None, None) + (hc, part, ss), (ref_hc, _, _) = out, ref + ratio = _hc_band_ratio(hc, ref_hc, block_out, inject, shared, gate, topk_weights, topk_rows) + assert ratio <= 1.0, f"hc_out {ratio:.4g}x the band" + _compare((hc, part, ss), op.reference(hc, None, None, gamma, first, c, n)) + + +def _build(spec, seed): + """The Triton entry's test operands at one row.""" + rows, hc_count, hidden, dtype = 1, HC_COUNT, HIDDEN, torch.bfloat16 + width = hc_count * hidden + n_cols = spec.get("n_cols", RANK + hc_count) + g = torch.Generator(device=DEV).manual_seed(seed) + + def normal(*shape, scale): + return torch.randn(*shape, generator=g, device=DEV) * scale + + hyper_input = normal(rows, width, scale=3.0).to(dtype) + if spec.get("layout") == "row_slice": + buf = normal(2 * rows, width, scale=1.0).to(dtype) + buf[::2] = hyper_input + hyper_input = buf[::2] + gamma = normal(width, scale=GAMMA_RMS) + gamma[0], gamma[width // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX + gamma = gamma.to(dtype) + n_rows = spec.get("n_rows", -(-n_cols // 128) * 128) + weight = torch.zeros(n_rows, width, device=DEV) + weight[:n_cols] = normal(n_cols, width, scale=DOWN_RMS) + weight[n_cols - hc_count : n_cols] *= INJECT_RMS / DOWN_RMS + first = weight.to(dtype).t() + block_out = inject = shared = gate_logits = None + if spec.get("inject", True): + block_out = normal(rows, hidden, scale=2.0).to(dtype) + inject = (2.0 * torch.sigmoid(normal(rows, hc_count, scale=1.0))).to(dtype) + if spec.get("moe", False): + shared = normal(rows, hidden, scale=2.0).to(dtype) + gate_logits = normal(rows, 1, scale=2.0).to(dtype) + if spec.get("gate") == "flat": + gate_logits = gate_logits.reshape(rows) + top_k = spec.get("fold", 0) + if not top_k: + return hyper_input, block_out, inject, gamma, first, hc_count, n_cols, shared, gate_logits + n_perm = rows * top_k + 37 + block_out = normal(n_perm, hidden, scale=2.0).to(dtype) + order = torch.randperm(n_perm, generator=g, device=DEV)[: rows * top_k] + topk_rows = order.reshape(rows, top_k).to(torch.int32) + if spec.get("drop"): + topk_rows[:, -3:] = -1 + topk_weights = torch.softmax(normal(rows, top_k, scale=1.0), -1).to(dtype) + return ( + hyper_input, + block_out, + inject, + gamma, + first, + hc_count, + n_cols, + shared, + gate_logits, + topk_weights, + topk_rows, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) + before = [a.clone() if isinstance(a, torch.Tensor) else a for a in args] + out = op(*args) + for a, b in zip(args, before): + assert not isinstance(a, torch.Tensor) or torch.equal(a, b), "every input is read-only" + if args[1] is None: + assert out[0] is args[0], "norm-only returns hyper_input itself" + if args[7] is None: + _compare(out, op.reference(*args)) + else: + _compare_moe(args, out, op.reference(*args)) + + +@pytest.mark.parametrize( + "spec", + [ + dict(), + dict(moe=True), + dict(moe=True, fold=10), + dict(moe=True, fold=10, drop=True), + dict(inject=False), + ], + ids=["inject", "combine", "fold", "fold-drop", "norm-only"], +) +def test_matches_the_triton_entry(spec) -> None: + """`hc_out` bitwise (same fp32 arithmetic and rounding); `part` and `ss` reduce in + another order, inside the band of the Triton entry's values.""" + args = _build(spec, seed=4242) + hc, part, ss = op(*args) + t_hc, t_part, t_ss = triton_op(*args) + assert torch.equal(hc, t_hc), "hc_out" + assert _band(part, t_part, PART_EPS) <= 1.0, "part" + assert torch.allclose(ss, t_ss, rtol=1e-5, atol=0), "ss" + print( + f"\n {spec}: part bitwise {(part == t_part).float().mean().item():.3f} " + f"({_band(part, t_part, PART_EPS):.3g}x the band apart)" + ) + + +def test_deterministic() -> None: + args = _build(dict(moe=True, fold=10), seed=3) + for a, b in zip(op(*args), op(*args)): + assert torch.equal(a, b) + + +def test_reads_after_the_wait_under_graph_replay() -> None: + """A producer kernel rewrites every activation inside the graph right before this op. + With PDL the op may start before the producer finishes; reading any activation before + `griddepcontrol.wait` would return the previous replay's value.""" + args = list(_build(dict(moe=True, fold=10), seed=8)) + sources = [a.clone() for a in (args[0], args[1], args[2], args[7], args[8], args[9])] + targets = (args[0], args[1], args[2], args[7], args[8], args[9]) + + def step(): + for dst, src in zip(targets, sources): + dst.copy_(src * 1.0) # an elementwise producer kernel per operand + return op(*args) + + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + for _ in range(2): + step() + torch.cuda.current_stream().wait_stream(warmup) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_out = step() + for trial in range(3): + fresh = _build(dict(moe=True, fold=10), seed=100 + trial) + for src, new in zip(sources, (fresh[0], fresh[1], fresh[2], fresh[7], fresh[8], fresh[9])): + src.copy_(new) + graph.replay() + torch.cuda.synchronize() + want = op(*[*args[:10], args[10]]) + for a, b in zip(graph_out, want): + assert torch.equal(a, b), f"replay {trial}" + + +def test_guards() -> None: + args = list(_build(dict(), seed=6)) + two = args[0].expand(2, -1).contiguous() + with validating(op), pytest.raises(AssertionError, match="one row"): + op(two, args[1].expand(2, -1).contiguous(), args[2].expand(2, -1).contiguous(), *args[3:]) + args[4] = args[4].contiguous() + with validating(op), pytest.raises(AssertionError): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix.py new file mode 100644 index 000000000000..306e1049b8d7 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix.py @@ -0,0 +1,157 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qwen4_exp_hc_decode_mix entry: cells inside +`qwen4_exp_gated_residual`'s band (`rtol = eps`, `atol = eps * max|ref|` of the +activation dtype) against the float64 reference for both `mixed` and +`inject_out`, wrong formulas outside, the stride guards, graph capture.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_decode_mix import ( # noqa: E501 + qwen4_exp_hc_decode_mix as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "qwen4_exp_hc_decode_mix requires a CUDA device" + +DEV = "cuda" +HC_COUNT, HIDDEN, RANK, EPS = 4, 2560, 320, 1e-6 +UP_RMS, GAMMA_RMS, GAMMA_ABSMAX = 0.0681, 0.477, 6.625 + + +def _band(got, want) -> float: + eps = torch.finfo(want.dtype).eps + w = want.double() + return ((got.double() - w).abs() / (eps * w.abs().max().item() + eps * w.abs())).max().item() + + +def _compare(out, ref) -> None: + (mixed, inj), (ref_mixed, ref_inj) = out, ref + assert (inj is None) == (ref_inj is None) + for got, want in [(mixed, ref_mixed)] + ([] if inj is None else [(inj, ref_inj)]): + assert got.dtype == want.dtype and got.shape == want.shape and got.is_contiguous() + if got.numel(): + assert _band(got, want) <= 1.0, f"{_band(got, want):.4g}x the band" + + +def _build(spec, seed): + rows = spec["rows"] + hc_count = spec.get("hc_count", HC_COUNT) + hidden = spec.get("hidden", HIDDEN) + rank = spec.get("rank", RANK) + dtype = spec.get("dtype", torch.bfloat16) + with_inject = spec.get("with_inject", True) + width = hc_count * hidden + g = torch.Generator(device=DEV).manual_seed(seed) + + def normal(*shape, scale): + return torch.randn(*shape, generator=g, device=DEV) * scale + + hc_out = normal(rows, width, scale=3.0).to(dtype) + if spec.get("layout") == "row_slice": + buf = normal(2 * rows, width, scale=1.0).to(dtype) + buf[::2] = hc_out + hc_out = buf[::2] + ss = hc_out.double().unflatten(-1, (hc_count, hidden)).pow(2).sum(-1).t().float().contiguous() + gamma = normal(width, scale=GAMMA_RMS) + gamma[0], gamma[width // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX + gamma = gamma.to(dtype) + # Per-stream partials sized so the combined logits reach this checkpoint's ~|67|. + n_pad = -(-(rank + hc_count) // 4) * 4 + part = normal(hc_count, rows, n_pad + 4, scale=60.0)[:, :, :n_pad] + up = normal(width, rank, scale=UP_RMS).to(dtype).t() + return part, ss, up, hc_out, gamma, EPS, hc_count, with_inject + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) + out = op(*args) + assert out[0].shape == (cell.spec["rows"], args[3].shape[1] // args[6]) + _compare(out, op.reference(*args)) + + +def test_measured_distance_from_float64() -> None: + for rows in (1, 16): + args = _build(dict(rows=rows), seed=rows) + (m, i), (rm, ri) = op(*args), op.reference(*args) + print(f"\n rows {rows}: mixed {_band(m, rm):.4f}x, inject {_band(i, ri):.4f}x the band") + + +def test_programs_agree_and_repeat() -> None: + args = _build(dict(rows=4), seed=3) + a, b = op(*args), op(*args) + assert torch.equal(a[0], b[0]) and torch.equal(a[1], b[1]) + + +def test_wrong_variants_land_outside_the_band() -> None: + part, ss, up, hc_out, gamma, eps, c, _ = args = _build(dict(rows=8), seed=7) + want_mixed, want_inj = op.reference(*args) + rstd = torch.rsqrt(ss.double().t() / HIDDEN + eps) + x = hc_out.double().unflatten(-1, (c, HIDDEN)) + g = (1.0 + gamma.double()).unflatten(-1, (c, HIDDEN)) + logits = torch.einsum("ts,stn->tn", rstd, part.double()) + raw_logits = part.double().sum(0) + + def mix(lg, normed, reduce="mean"): + gate = torch.sigmoid(torch.nn.functional.silu(lg[:, :RANK] / c) @ up.double()) + v = gate.unflatten(-1, (c, HIDDEN)) * normed + return v.mean(-2) if reduce == "mean" else v.sum(-2) + + normed = x * rstd.unsqueeze(-1) * g + variants = { + "sum instead of mean": mix(logits, normed, "sum"), + "gates the raw stream": mix(logits, x), + "w instead of 1 + w": mix(logits, x * rstd.unsqueeze(-1) * (g - 1.0)), + "no rstd on the logits": mix(raw_logits, normed), + "no /hc_count before silu": mix(logits * c, normed), + } + for name, value in variants.items(): + ratio = _band(value.to(torch.bfloat16), want_mixed) + print(f"\n mixed {name:<28} {ratio:10.4g}x the band") + assert ratio > 5.0, name + zi = logits[:, RANK : RANK + c] + for name, value in { + "inject missing the factor of 2": torch.sigmoid(zi / c), + "inject without rstd": 2.0 * torch.sigmoid(raw_logits[:, RANK : RANK + c] / c), + }.items(): + ratio = _band(value.to(torch.bfloat16), want_inj) + print(f" {name:<34} {ratio:10.4g}x the band") + assert ratio > 5.0, name + + +def test_guard_refuses_a_strided_ss() -> None: + """`ss` laid out `[T, C]` and viewed as `[C, T]` is read row-major and + answers plausibly wrong unguarded; the guard refuses it.""" + args = list(_build(dict(rows=8), seed=6)) + want = op.reference(*args)[0] + args[1] = args[1].t().contiguous().t() + assert _band(op(*args)[0], want) > 1.0, "a transposed ss should answer outside the band" + with validating(op), pytest.raises(AssertionError): + op(*args) + + +def test_cuda_graph_capture_and_replay() -> None: + args = list(_build(dict(rows=1), seed=8)) + 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): + graph_out = op(*args) + fresh = _build(dict(rows=1), seed=18) + for k in (0, 1, 3): + args[k].copy_(fresh[k]) + graph.replay() + torch.cuda.synchronize() + eager = op(*args) + assert torch.equal(graph_out[0], eager[0]) and torch.equal(graph_out[1], eager[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_hc_decode_mix_cute.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_cute.py new file mode 100644 index 000000000000..fc1a21f55ba8 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_cute.py @@ -0,0 +1,132 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qwen4_exp_hc_decode_mix_cute entry (the plain build of the CuTe mix +kernel): cells against the Triton entry's float64 reference in its band (`mixed` and +`inject_out` in `rtol = eps`, `atol = eps * max|ref|`), agreement with the Triton entry +inside the same band, a stale-read check under CUDA-graph replay with PDL, 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.hc.qwen4_exp_hc_decode_mix import ( # noqa: E501 + qwen4_exp_hc_decode_mix as triton_op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_decode_mix_fp4_cute import ( # noqa: E501 + qwen4_exp_hc_decode_mix_cute as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "qwen4_exp_hc_decode_mix_cute requires a CUDA device" + +DEV = "cuda" +HC_COUNT, HIDDEN, RANK, EPS = 4, 2560, 320, 1e-6 +UP_RMS, GAMMA_RMS, GAMMA_ABSMAX = 0.0681, 0.477, 6.625 + + +def _band(got, want) -> float: + eps = torch.finfo(want.dtype).eps + w = want.double() + return ((got.double() - w).abs() / (eps * w.abs().max().item() + eps * w.abs())).max().item() + + +def _compare(out, ref) -> None: + (mixed, inj), (ref_mixed, ref_inj) = out, ref + assert (inj is None) == (ref_inj is None) + for got, want in [(mixed, ref_mixed)] + ([] if inj is None else [(inj, ref_inj)]): + assert got.dtype == want.dtype and got.shape == want.shape and got.is_contiguous() + assert _band(got, want) <= 1.0, f"{_band(got, want):.4g}x the band" + + +def _build(spec, seed): + rows, hc_count, hidden, rank, dtype = 1, HC_COUNT, HIDDEN, RANK, torch.bfloat16 + with_inject = spec.get("with_inject", True) + width = hc_count * hidden + g = torch.Generator(device=DEV).manual_seed(seed) + + def normal(*shape, scale): + return torch.randn(*shape, generator=g, device=DEV) * scale + + hc_out = normal(rows, width, scale=3.0).to(dtype) + if spec.get("layout") == "row_slice": + buf = normal(2 * rows, width, scale=1.0).to(dtype) + buf[::2] = hc_out + hc_out = buf[::2] + ss = hc_out.double().unflatten(-1, (hc_count, hidden)).pow(2).sum(-1).t().float().contiguous() + gamma = normal(width, scale=GAMMA_RMS) + gamma[0], gamma[width // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX + gamma = gamma.to(dtype) + n_pad = -(-(rank + hc_count) // 4) * 4 + part = normal(hc_count, rows, n_pad + 4, scale=60.0)[:, :, :n_pad] + up = normal(width, rank, scale=UP_RMS).to(dtype).t() + return [part, ss, up, hc_out, gamma, EPS, hc_count, with_inject] + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) + out = op(*args) + assert out[0].shape == (1, HIDDEN) + _compare(out, op.reference(*args)) + + +@pytest.mark.parametrize("with_inject", [True, False]) +def test_matches_the_triton_entry(with_inject) -> None: + args = _build(dict(with_inject=with_inject), seed=77) + mixed, inj = op(*args) + t_mixed, t_inj = triton_op(*args) + assert _band(mixed, t_mixed) <= 1.0 + if with_inject: + assert _band(inj, t_inj) <= 1.0 + print(f"\n mixed bitwise the Triton entry's: {(mixed == t_mixed).float().mean().item():.4f}") + + +def test_deterministic() -> None: + args = _build(dict(), seed=3) + for a, b in zip(op(*args), op(*args)): + assert torch.equal(a, b) + + +def test_reads_after_the_wait_under_graph_replay() -> None: + """Producer kernels rewrite `part`, `ss` and `hc_out` inside the graph right before + the op; with PDL a read before `griddepcontrol.wait` would see the previous replay's.""" + args = _build(dict(), seed=8) + targets = (args[0], args[1], args[3]) + sources = [t.clone() for t in targets] + + def step(): + for dst, src in zip(targets, sources): + dst.copy_(src * 1.0) + return op(*args) + + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + for _ in range(2): + step() + torch.cuda.current_stream().wait_stream(warmup) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_out = step() + for trial in range(3): + fresh = _build(dict(), seed=100 + trial) + for src, new in zip(sources, (fresh[0], fresh[1], fresh[3])): + src.copy_(new) + graph.replay() + torch.cuda.synchronize() + for a, b in zip(graph_out, op(*args)): + assert torch.equal(a, b), f"replay {trial}" + + +def test_guards() -> None: + args = _build(dict(), seed=6) + two = args[3].expand(2, -1).contiguous() + with validating(op), pytest.raises(AssertionError): + op(args[0].expand(-1, 2, -1), args[1].expand(-1, 2), args[2], two, *args[4:]) + contiguous_up = args[2].contiguous() + with validating(op), pytest.raises(AssertionError): + op(args[0], args[1], contiguous_up, *args[3:]) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4.py new file mode 100644 index 000000000000..d760d66b80ab --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4.py @@ -0,0 +1,190 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qwen4_exp_hc_decode_mix_fp4 entry: `mixed` and `inject_out` +inside `qwen4_exp_hc_decode_mix`'s band against the float64 reference, the +NVFP4 outputs exact on the returned `mixed` (outside the reciprocal tie window) +and bitwise against `qwen4_exp_hc_decode_mix` + `torch.ops.trtllm.fp4_quantize`, +guards, graph capture.""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_decode_mix import ( # noqa: E501 + qwen4_exp_hc_decode_mix, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_decode_mix_fp4 import ( # noqa: E501 + nvfp4_near_tie, + nvfp4_reference, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_decode_mix_fp4 import ( # noqa: E501 + qwen4_exp_hc_decode_mix_fp4 as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "qwen4_exp_hc_decode_mix_fp4 requires a CUDA device" + +DEV = "cuda" +HC_COUNT, HIDDEN, RANK, EPS = 4, 2560, 320, 1e-6 +UP_RMS, GAMMA_RMS, GAMMA_ABSMAX = 0.0681, 0.477, 6.625 + + +def _band(got, want) -> float: + eps = torch.finfo(want.dtype).eps + w = want.double() + return ((got.double() - w).abs() / (eps * w.abs().max().item() + eps * w.abs())).max().item() + + +def _unpack(xq: torch.Tensor) -> torch.Tensor: + return torch.stack([xq & 0xF, xq >> 4], -1).flatten(-2) + + +def _check_quant(mixed, xq, xsf, gs) -> None: + want_q, want_sf = nvfp4_reference(mixed, gs) + assert xq.dtype == torch.uint8 and xq.shape == want_q.shape and xq.is_contiguous() + assert xsf.dtype == torch.uint8 and xsf.shape == want_sf.shape + assert torch.equal(xsf, want_sf), f"{(xsf != want_sf).sum().item()} scale bytes differ" + stray = (_unpack(xq) != _unpack(want_q)) & ~nvfp4_near_tie(mixed, gs) + assert not stray.any(), f"{stray.sum().item()} e2m1 codes differ outside the tie window" + + +def _compare(out, ref, gs) -> None: + """`mixed` / `inject_out` in `rtol = eps`, `atol = eps * max|ref|` (the band + `qwen4_exp_hc_decode_mix` is certified in); NVFP4 exact on the returned + `mixed`, and equal to the reference's on 16-blocks whose `mixed` matches.""" + (mixed, inj, xq, xsf), (ref_mixed, ref_inj, ref_q, ref_sf) = out, ref + assert (inj is None) == (ref_inj is None) + for got, want in [(mixed, ref_mixed)] + ([] if inj is None else [(inj, ref_inj)]): + assert got.dtype == want.dtype and got.shape == want.shape and got.is_contiguous() + if got.numel(): + assert _band(got, want) <= 1.0, f"{_band(got, want):.4g}x the band" + if mixed.numel() == 0: + assert xq.numel() == 0 and xsf.numel() == 0 + return + _check_quant(mixed, xq, xsf, gs) + rows = mixed.shape[0] + same = (mixed.view(rows, -1, 16) == ref_mixed.view(rows, -1, 16)).all(-1) + assert torch.equal(xsf.view(rows, -1)[same], ref_sf.view(rows, -1)[same]) + q, rq = _unpack(xq).view(rows, -1, 16), _unpack(ref_q).view(rows, -1, 16) + tie = nvfp4_near_tie(ref_mixed, gs).view(rows, -1, 16) + assert not ((q != rq) & ~tie)[same].any() + + +def _build(spec, seed): + rows = spec["rows"] + hc_count = spec.get("hc_count", HC_COUNT) + hidden = spec.get("hidden", HIDDEN) + rank = spec.get("rank", RANK) + dtype = spec.get("dtype", torch.bfloat16) + with_inject = spec.get("with_inject", True) + width = hc_count * hidden + g = torch.Generator(device=DEV).manual_seed(seed) + + def normal(*shape, scale): + return torch.randn(*shape, generator=g, device=DEV) * scale + + hc_out = normal(rows, width, scale=3.0).to(dtype) + if spec.get("layout") == "row_slice": + buf = normal(2 * rows, width, scale=1.0).to(dtype) + buf[::2] = hc_out + hc_out = buf[::2] + ss = hc_out.double().unflatten(-1, (hc_count, hidden)).pow(2).sum(-1).t().float().contiguous() + gamma = normal(width, scale=GAMMA_RMS) + gamma[0], gamma[width // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX + gamma = gamma.to(dtype) + n_pad = -(-(rank + hc_count) // 4) * 4 + part = normal(hc_count, rows, n_pad + 4, scale=60.0)[:, :, :n_pad] + up = normal(width, rank, scale=UP_RMS).to(dtype).t() + # A calibrated activation global scale: 448 * 6 / amax of the mix (~|8|). + gs = torch.tensor([448.0 * 6.0 / 8.0], device=DEV, dtype=torch.float32) + return part, ss, up, hc_out, gamma, EPS, hc_count, with_inject, gs + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) + out = op(*args) + assert out[0].shape == (cell.spec["rows"], args[3].shape[1] // args[6]) + _compare(out, op.reference(*args), args[8]) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_bitwise_the_composition_it_replaces(cell) -> None: + """`qwen4_exp_hc_decode_mix` then `torch.ops.trtllm.fp4_quantize` (linear + scales). The NVFP4 outputs are bitwise `fp4_quantize` of the returned + `mixed`, ties included; `inject_out` is bitwise `qwen4_exp_hc_decode_mix`'s. + `mixed` is held to `qwen4_exp_hc_decode_mix`'s within the band: the + 16-column program tile reduces the up projection in another order than the + 8-column one, so a few elements in ten thousand round the other way + (measured: 0..5 of 10240 elements at 1..4 rows).""" + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) + mixed, inj, xq, xsf = op(*args) + old_mixed, old_inj = qwen4_exp_hc_decode_mix(*args[:8]) + assert (inj is None and old_inj is None) or torch.equal(inj, old_inj) + if mixed.numel(): + diff = (mixed != old_mixed).sum().item() + print( + f"\n {mixed.dtype}: {diff} elements differ from decode_mix, " + f"{_band(mixed, old_mixed):.4f}x the band" + ) + assert _band(mixed, old_mixed) <= 1.0 + if mixed.shape[0] == 0: + return + old_q, old_sf = torch.ops.trtllm.fp4_quantize(mixed, args[8], 16, False, False) + assert torch.equal(xq, old_q.view(torch.uint8)) + assert torch.equal(xsf, old_sf.view(torch.uint8).reshape(-1)) + + +def test_wrong_recipes_fail_the_quant_check() -> None: + """Each plausible wrong layout or recipe is refused by `_check_quant`.""" + args = _build(dict(rows=16), seed=3) + mixed, _, xq, xsf = op(*args) + gs = args[8] + _check_quant(mixed, xq, xsf, gs) + swapped = ((xq & 0xF) << 4) | (xq >> 4) + shifted = torch.roll(xsf.view(16, -1), 1, dims=1).reshape(-1) + half_gs = nvfp4_reference(mixed, gs * 0.5) + for name, bad in ( + ("nibbles swapped", (mixed, swapped, xsf)), + ("scales one block off", (mixed, xq, shifted)), + ("global scale halved", (mixed,) + half_gs), + ): + with pytest.raises(AssertionError): + _check_quant(*bad, gs) + print(f"\n {name:<24} refused") + + +def test_guards() -> None: + args = list(_build(dict(rows=4), seed=4)) + with validating(op), pytest.raises(AssertionError): + op(*args[:8], torch.cat([args[8], args[8]])) + bad = list(_build(dict(rows=2, hidden=136, rank=32), seed=5)) + with validating(op), pytest.raises(AssertionError): + op(*bad) + + +def test_cuda_graph_capture_and_replay() -> None: + args = list(_build(dict(rows=1), seed=8)) + 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): + graph_out = op(*args) + fresh = _build(dict(rows=1), seed=18) + for k in (0, 1, 3): + args[k].copy_(fresh[k]) + args[8].mul_(0.5) + graph.replay() + torch.cuda.synchronize() + eager = op(*args) + for got, want in zip(graph_out, eager): + assert torch.equal(got, want) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py new file mode 100644 index 000000000000..29724ea6a96a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py @@ -0,0 +1,157 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qwen4_exp_hc_decode_mix_fp4_cute entry: cells against the Triton +entry's float64 reference in its band (`mixed` and `inject_out` in `rtol = eps`, +`atol = eps * max|ref|`; the NVFP4 exact on the returned `mixed` outside the +reciprocal tie window), agreement with the Triton entry inside the same band, a +stale-read check under CUDA-graph replay with PDL, 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.hc.qwen4_exp_hc_decode_mix_fp4 import ( # noqa: E501 + nvfp4_near_tie, + nvfp4_reference, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_decode_mix_fp4 import ( # noqa: E501 + qwen4_exp_hc_decode_mix_fp4 as triton_op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_decode_mix_fp4_cute import ( # noqa: E501 + qwen4_exp_hc_decode_mix_fp4_cute as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "qwen4_exp_hc_decode_mix_fp4_cute requires a CUDA device" + +DEV = "cuda" +HC_COUNT, HIDDEN, RANK, EPS = 4, 2560, 320, 1e-6 +UP_RMS, GAMMA_RMS, GAMMA_ABSMAX = 0.0681, 0.477, 6.625 + + +def _band(got, want) -> float: + eps = torch.finfo(want.dtype).eps + w = want.double() + return ((got.double() - w).abs() / (eps * w.abs().max().item() + eps * w.abs())).max().item() + + +def _compare(out, ref) -> None: + (mixed, inj), (ref_mixed, ref_inj) = out[:2], ref[:2] + assert (inj is None) == (ref_inj is None) + for got, want in [(mixed, ref_mixed)] + ([] if inj is None else [(inj, ref_inj)]): + assert got.dtype == want.dtype and got.shape == want.shape and got.is_contiguous() + assert _band(got, want) <= 1.0, f"{_band(got, want):.4g}x the band" + + +def _unpack(xq): + return torch.stack([xq & 0xF, xq >> 4], -1).flatten(-2) + + +def _check_quant(mixed, xq, xsf, gs) -> None: + """The NVFP4 outputs exact on the returned `mixed`: scales bitwise, e2m1 codes + bitwise outside the window where the approximate reciprocals may pick either side.""" + want_q, want_sf = nvfp4_reference(mixed, gs) + assert xq.dtype == torch.uint8 and xq.shape == want_q.shape and xq.is_contiguous() + assert xsf.dtype == torch.uint8 and xsf.shape == want_sf.shape + assert torch.equal(xsf, want_sf), f"{(xsf != want_sf).sum().item()} scale bytes differ" + stray = (_unpack(xq) != _unpack(want_q)) & ~nvfp4_near_tie(mixed, gs) + assert not stray.any(), f"{stray.sum().item()} e2m1 codes differ outside the tie window" + + +def _build(spec, seed): + rows, hc_count, hidden, rank, dtype = 1, HC_COUNT, HIDDEN, RANK, torch.bfloat16 + with_inject = spec.get("with_inject", True) + width = hc_count * hidden + g = torch.Generator(device=DEV).manual_seed(seed) + + def normal(*shape, scale): + return torch.randn(*shape, generator=g, device=DEV) * scale + + hc_out = normal(rows, width, scale=3.0).to(dtype) + if spec.get("layout") == "row_slice": + buf = normal(2 * rows, width, scale=1.0).to(dtype) + buf[::2] = hc_out + hc_out = buf[::2] + ss = hc_out.double().unflatten(-1, (hc_count, hidden)).pow(2).sum(-1).t().float().contiguous() + gamma = normal(width, scale=GAMMA_RMS) + gamma[0], gamma[width // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX + gamma = gamma.to(dtype) + n_pad = -(-(rank + hc_count) // 4) * 4 + part = normal(hc_count, rows, n_pad + 4, scale=60.0)[:, :, :n_pad] + up = normal(width, rank, scale=UP_RMS).to(dtype).t() + # A calibrated activation global scale: 448 * 6 / amax of the mix (~|8|). + gs = torch.tensor([448.0 * 6.0 / 8.0], device=DEV, dtype=torch.float32) + return [part, ss, up, hc_out, gamma, EPS, hc_count, with_inject, gs] + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) + out = op(*args) + assert out[0].shape == (1, HIDDEN) + _compare(out, op.reference(*args)) + _check_quant(out[0], out[2], out[3], args[8]) + + +@pytest.mark.parametrize("with_inject", [True, False]) +def test_matches_the_triton_entry(with_inject) -> None: + """`mixed` / `inject_out` inside the band of the Triton entry's; the NVFP4 of each + is exact on its own `mixed`, so where the two `mixed` agree the codes agree.""" + args = _build(dict(with_inject=with_inject), seed=77) + mixed, inj, _, _ = op(*args) + t_mixed, t_inj, _, _ = triton_op(*args) + assert _band(mixed, t_mixed) <= 1.0 + if with_inject: + assert _band(inj, t_inj) <= 1.0 + print(f"\n mixed bitwise the Triton entry's: {(mixed == t_mixed).float().mean().item():.4f}") + + +def test_deterministic() -> None: + args = _build(dict(), seed=3) + for a, b in zip(op(*args), op(*args)): + assert torch.equal(a, b) + + +def test_reads_after_the_wait_under_graph_replay() -> None: + """Producer kernels rewrite `part`, `ss` and `hc_out` inside the graph right before + the op; with PDL a read before `griddepcontrol.wait` would see the previous replay's.""" + args = _build(dict(), seed=8) + targets = (args[0], args[1], args[3]) + sources = [t.clone() for t in targets] + + def step(): + for dst, src in zip(targets, sources): + dst.copy_(src * 1.0) + return op(*args) + + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + for _ in range(2): + step() + torch.cuda.current_stream().wait_stream(warmup) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_out = step() + for trial in range(3): + fresh = _build(dict(), seed=100 + trial) + for src, new in zip(sources, (fresh[0], fresh[1], fresh[3])): + src.copy_(new) + graph.replay() + torch.cuda.synchronize() + for a, b in zip(graph_out, op(*args)): + assert torch.equal(a, b), f"replay {trial}" + + +def test_guards() -> None: + args = _build(dict(), seed=6) + two = args[3].expand(2, -1).contiguous() + with validating(op), pytest.raises(AssertionError): + op(args[0].expand(-1, 2, -1), args[1].expand(-1, 2), args[2], two, *args[4:]) + contiguous_up = args[2].contiguous() + with validating(op), pytest.raises(AssertionError): + op(args[0], args[1], contiguous_up, *args[3:]) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first.py new file mode 100644 index 000000000000..8ba9eb3d73e6 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first.py @@ -0,0 +1,186 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qwen4_exp_hc_decode_tc_first entry: cells in the band against +the float64 reference, wrong formulas outside, determinism, the guard, graph +capture, and the decode mixer built on it against the gated-residual float64 +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.hc.qwen4_exp_hc_decode_tc_first import ( # noqa: E501 + qwen4_exp_hc_decode_tc_first as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "qwen4_exp_hc_decode_tc_first requires a CUDA device" + +DEV = "cuda" +RANK = 320 + + +def _ratio(out, ref) -> float: + """Worst element's distance from `ref` as a fraction of the band: one ulp of + the element (`eps * |ref|`) plus `eps / 16` of the tensor's largest value. + An fp32-accumulated 10240-term projection rounded once can sit one ulp off + the float64 value at a rounding midpoint; the absolute term covers SiLU's + cancellation near zero, where an fp32 logit error of ~1e-6 relative to the + row's scale is many ulps of a tiny output.""" + eps = torch.finfo(ref.dtype).eps + r = ref.double() + return ((out.double() - r).abs() / (eps * r.abs() + eps / 16 * r.abs().max())).max().item() + + +def _compare(out, ref) -> None: + for got, want in zip(out, ref): + assert (got is None) == (want is None) + if got is None: + continue + assert got.dtype == want.dtype and got.shape == want.shape and got.is_contiguous() + if got.numel(): + assert _ratio(got, want) <= 1.0, f"{_ratio(got, want):.4g}x the band" + + +def _build(spec, seed): + hc, hidden = spec.get("hc_count", 4), spec.get("hidden", 2560) + rows, width, inject = spec["rows"], hc * hidden, spec.get("inject", True) + g = torch.Generator(device=DEV).manual_seed(seed) + # This checkpoint's layer statistics: down ~N(0, 0.0296), block_inject ~N(0, 0.0107), + # the normed stream unit-RMS per stream times (1 + gamma), gamma ~N(0, 0.477). + down = torch.randn(RANK, width, generator=g, device=DEV) * 0.0296 + parts = [down] + ([torch.randn(hc, width, generator=g, device=DEV) * 0.0107] if inject else []) + w = torch.cat(parts) + n_rows = spec.get("n_rows", 384) + if n_rows > w.shape[0]: + w = torch.cat([w, torch.zeros(n_rows - w.shape[0], width, device=DEV)]) + gamma = torch.randn(width, generator=g, device=DEV) * 0.477 + x = torch.randn(rows, hc, hidden, generator=g, device=DEV) + x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) + normed = (x.flatten(-2) * (1.0 + gamma)).to(torch.bfloat16) + return normed, w.to(torch.bfloat16).t(), RANK, hc, inject + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) + low, inject = op(*args) + assert low.shape == (cell.spec["rows"], RANK) + _compare((low, inject), op.reference(*args)) + + +def test_deterministic() -> None: + args = _build(dict(rows=128), seed=7) + a, b = op(*args), op(*args) + assert torch.equal(a[0], b[0]) and torch.equal(a[1], b[1]) + + +def test_wrong_variants_land_outside_the_band() -> None: + normed, first, rank, c, wi = _build(dict(rows=64), seed=3) + want_low, want_inject = op.reference(normed, first, rank, c, wi) + got = op(normed, first, rank, c, wi) + lg = normed.double() @ first.double() + for name, value, want in ( + ("silu without /hc_count", F.silu(lg[:, :rank]), want_low), + ("sigmoid instead of silu", torch.sigmoid(lg[:, :rank] / c), want_low), + ("columns shifted by one", F.silu(lg[:, 1 : rank + 1] / c), want_low), + ("inject missing the 2", torch.sigmoid(lg[:, rank : rank + c] / c), want_inject), + ( + "inject from the padding", + 2 * torch.sigmoid(lg[:, rank + c : rank + 2 * c] / c), + want_inject, + ), + ): + ratio = _ratio(value.to(want.dtype), want) + print(f"\n {name:<28} {ratio:10.4g}x the band (op: {_ratio(got[0], want_low):.3g}x)") + assert ratio > 5.0, name + + +def test_guard_refuses_a_contiguous_weight() -> None: + normed, first, rank, c, wi = _build(dict(rows=4), seed=4) + with validating(op), pytest.raises(AssertionError): + op(normed, first.contiguous(), rank, c, wi) + + +def test_guard_refuses_more_than_one_tile_of_rows() -> None: + normed, first, rank, c, wi = _build(dict(rows=129), seed=4) + with validating(op), pytest.raises(AssertionError): + op(normed, first, rank, c, wi) + + +def test_cuda_graph_capture_and_replay() -> None: + normed, first, rank, c, wi = _build(dict(rows=8), seed=5) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + op(normed, first, rank, c, wi) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + g_low, g_inject = op(normed, first, rank, c, wi) + fresh = _build(dict(rows=8), seed=6)[0] + normed.copy_(fresh) + graph.replay() + torch.cuda.synchronize() + low, inject = op(fresh, first, rank, c, wi) + assert torch.equal(g_low, low) and torch.equal(g_inject, inject) + + +@pytest.mark.parametrize("rows", [2, 33, 128]) +@pytest.mark.parametrize("with_inject", [True, False]) +def test_mixer_matches_the_gated_residual_reference(rows, with_inject) -> None: + """The decode mixer as `DecodeTarget._hc_mixer` runs it from 2 rows -- this op, + then `qwen4_exp_hc_up_gate_mean` -- against `qwen4_exp_gated_residual.reference` + (float64 from `normed`) in the band that entry certifies its own op in, beside + the cuBLAS + `qwen4_exp_hc_lowrank_act` mixer it replaces.""" + from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm + from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501 + gate_ratio, + qwen4_exp_gated_residual, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_lowrank_act import ( # noqa: E501 + qwen4_exp_hc_lowrank_act, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_up_gate_mean import ( # noqa: E501 + qwen4_exp_hc_up_gate_mean, + ) + + hc, hidden, eps = 4, 2560, 1e-6 + width = hc * hidden + g = torch.Generator(device=DEV).manual_seed(rows) + + def normal(*shape, scale): + return (torch.randn(*shape, generator=g, device=DEV) * scale).to(torch.bfloat16) + + # This checkpoint's layer-0 hyper-connection element statistics. + down, up = normal(RANK, width, scale=0.0296), normal(width, RANK, scale=0.0681) + inject_w, gamma = normal(hc, width, scale=0.0107), normal(width, scale=0.477) + gamma[0], gamma[width // 2] = 6.625, -6.625 + x = torch.randn(rows, hc, hidden, generator=g, device=DEV) + x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) + normed = (x.flatten(-2) * (1.0 + gamma.float())).to(torch.bfloat16) + first = torch.cat([down, inject_w], 0) if with_inject else down + first = torch.cat( + [first, torch.zeros(384 - first.shape[0], width, device=DEV, dtype=first.dtype)] + ) + low, inject = op(normed, first.t(), RANK, hc, with_inject) + mixed = qwen4_exp_hc_up_gate_mean(low, up.t(), normed, hc) + want = qwen4_exp_gated_residual.reference( + normed, down, up, inject_w if with_inject else None, hc + ) + assert gate_ratio(mixed, want[0]) <= 1.0, f"mixed {gate_ratio(mixed, want[0]):.3f}x the band" + if with_inject: + assert gate_ratio(inject, want[1]) <= 1.0, f"inject {gate_ratio(inject, want[1]):.3f}x" + old_low, _ = qwen4_exp_hc_lowrank_act( + cublas_mm(normed, first.t(), None, torch.float32), RANK, hc, with_inject, torch.bfloat16 + ) + prev = qwen4_exp_hc_up_gate_mean(old_low, up.t(), normed, hc) + print( + f"\n rows {rows}: tc first {gate_ratio(mixed, want[0]):.3f}x band, " + f"cublas + lowrank_act {gate_ratio(prev, want[0]):.3f}x" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_gate_mean.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_gate_mean.py new file mode 100644 index 000000000000..a6f4b23d4df7 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_gate_mean.py @@ -0,0 +1,145 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qwen4_exp_hc_gate_mean entry: cells in the band +`qwen4_exp_gated_residual`'s test derives for `mixed`, wrong formulas outside, +graph capture.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501 + gate_ratio, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_gate_mean import ( # noqa: E501 + qwen4_exp_hc_gate_mean as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "qwen4_exp_hc_gate_mean requires a CUDA device" + +DEV = "cuda" + + +def _compare(out, ref) -> None: + """`rtol = eps`, `atol = eps * max|ref|` (`gate_ratio <= 1`): `mixed` is a + mean of signed products, so near-zero elements are cancellations whose + error is set by the stream's scale -- the band `qwen4_exp_gated_residual` + derives for the same output.""" + assert out.dtype == ref.dtype and out.shape == ref.shape and out.is_contiguous() + if out.numel(): + assert gate_ratio(out, ref) <= 1.0, f"{gate_ratio(out, ref):.4g}x the band" + + +def _build(spec, seed): + hc_count, hidden = spec.get("hc_count", 4), spec.get("hidden", 2560) + dtype = spec.get("dtype", torch.bfloat16) + g = torch.Generator(device=DEV).manual_seed(seed) + rows, width = spec["rows"], hc_count * hidden + # Gate logits of this checkpoint's up projection span about +-8; the normed stream +-30. + gate_logits = (torch.randn(rows, width, generator=g, device=DEV) * 3).to(dtype) + normed = (torch.randn(rows, width, generator=g, device=DEV) * 6).to(dtype) + return gate_logits, normed, 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 — fixture by name + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) + mixed = op(*args) + assert mixed.shape == (cell.spec["rows"], cell.spec.get("hidden", 2560)) + _compare(mixed, op.reference(*args)) + + +def test_wrong_variants_land_outside_the_band() -> None: + gate_logits, normed, c = _build(dict(rows=128), seed=3) + want = op.reference(gate_logits, normed, c) + h = normed.shape[1] // c + g = torch.sigmoid(gate_logits.double()).unflatten(-1, (c, h)) + x = normed.double().unflatten(-1, (c, h)) + for name, value in ( + ("sum instead of mean", (g * x).sum(-2)), + ("no sigmoid", (gate_logits.double().unflatten(-1, (c, h)) * x).mean(-2)), + ("stream-minor packing", (g.transpose(-1, -2).reshape(-1, c, h) * x).mean(-2)), + ("ungated mean", x.mean(-2)), + ): + ratio = gate_ratio(value.to(want.dtype), want) + print(f"\n {name:<24} {ratio:10.4g}x the band") + assert ratio > 5.0, name + + +def test_guard_refuses_mismatched_shapes() -> None: + gate_logits, normed, c = _build(dict(rows=4), seed=4) + with validating(op), pytest.raises(AssertionError): + op(gate_logits[:2], normed, c) + + +def test_cuda_graph_capture_and_replay() -> None: + gate_logits, normed, c = _build(dict(rows=8), seed=5) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + op(gate_logits, normed, c) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + g_mixed = op(gate_logits, normed, c) + fresh = _build(dict(rows=8), seed=6) + gate_logits.copy_(fresh[0]) + normed.copy_(fresh[1]) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(g_mixed, op(fresh[0], fresh[1], c)) + + +@pytest.mark.parametrize("rows", [1, 33, 128, 2048, 8192]) +@pytest.mark.parametrize("with_inject", [True, False]) +def test_mixer_matches_the_gated_residual_reference(rows, with_inject) -> None: + """The mixer as the forward runs it -- `cublas_mm` over the zero-padded + `[down; inject]` (fp32 out), `qwen4_exp_hc_lowrank_act`, `cublas_mm` for + the up projection, this op -- against `qwen4_exp_gated_residual.reference` + (float64 from `normed`), in the band that entry certifies its own op in.""" + from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm + from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501 + qwen4_exp_gated_residual, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_lowrank_act import ( # noqa: E501 + qwen4_exp_hc_lowrank_act, + ) + + hc, hidden, rank, eps = 4, 2560, 320, 1e-6 + width = hc * hidden + g = torch.Generator(device=DEV).manual_seed(rows) + + def normal(*shape, scale): + return (torch.randn(*shape, generator=g, device=DEV) * scale).to(torch.bfloat16) + + # This checkpoint's layer-0 hyper-connection element statistics. + down, up = normal(rank, width, scale=0.0296), normal(width, rank, scale=0.0681) + inject_w, gamma = normal(hc, width, scale=0.0107), normal(width, scale=0.477) + gamma[0], gamma[width // 2] = 6.625, -6.625 + x = torch.randn(rows, hc, hidden, generator=g, device=DEV) + x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) + normed = (x.flatten(-2) * (1.0 + gamma.float())).to(torch.bfloat16) + first = torch.cat([down, inject_w], 0) if with_inject else down + first = torch.cat( + [first, torch.zeros(384 - first.shape[0], width, device=DEV, dtype=first.dtype)] + ) + low, inject = qwen4_exp_hc_lowrank_act( + cublas_mm(normed, first.t(), None, torch.float32), rank, hc, with_inject, torch.bfloat16 + ) + mixed = op(cublas_mm(low, up.t()), normed, hc) + want = qwen4_exp_gated_residual.reference( + normed, down, up, inject_w if with_inject else None, hc + ) + _compare(mixed, want[0]) + if with_inject: + _compare(inject, want[1]) + old = qwen4_exp_gated_residual(normed, down, up, inject_w if with_inject else None, hc) + print( + f"\n rows {rows}: new mixed {gate_ratio(mixed, want[0]):.3f}x band, " + f"old op {gate_ratio(old[0], want[0]):.3f}x" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_inject_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_inject_norm.py new file mode 100644 index 000000000000..aa6de075afca --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_inject_norm.py @@ -0,0 +1,422 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qwen4_exp_hc_inject_norm entry. + +Cells are gated at one ulp of the activation dtype against the float64 +reference; the MoE-combine cells widen the residual's band by the combine's own +rounding (see `_compare_moe`). Beyond them: agreement with the compositions it +replaces, the separation from plausible wrong formulas, the norm-only aliasing +claim, and graph capture. +""" + +import sys +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_inject_norm import ( # noqa: E501 + qwen4_exp_hc_inject_norm as op, +) +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, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp4_block_scale_moe_renorm_permuted import ( # noqa: E501 + finalize, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fused_sigmoid_gate_mul_add import ( # noqa: E501 + fused_sigmoid_gate_mul_add, +) +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, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "qwen4_exp_hc_inject_norm requires a CUDA device" + +#: The op's module (its kernels and dispatch guards). +M = sys.modules[type(op).__module__] + +DEV = "cuda" +HC_COUNT = 4 +HIDDEN = 2560 +EPS = 1e-6 +# This checkpoint's grouped-norm weight statistics (layer-0 hyper-connection norm). +GAMMA_RMS = 0.477 +GAMMA_ABSMAX = 6.625 + + +def _ulp_ratio(out: torch.Tensor, ref: torch.Tensor) -> float: + """Worst element's distance from `ref` in units of `eps(dtype) * |ref|` + (one ulp of the reference value, or a little more). 1 means one ulp.""" + eps = torch.finfo(ref.dtype).eps + ref64 = ref.double() + tiny = torch.finfo(ref.dtype).tiny + return ((out.double() - ref64).abs() / (eps * ref64.abs() + tiny)).max().item() + + +def _compare(out, ref) -> None: + """Both outputs within one ulp of the float64 result rounded once. + + The op rounds each output once from an fp32 value whose own error is + ~1e-6 relative (a 16-bit product is exact in fp32; the RMS is a 2560-term + fp32 sum), so it can differ from the once-rounded float64 value only where + the two straddle a rounding midpoint: by exactly one ulp.""" + (hc, normed), (ref_hc, ref_normed) = out, ref + for got, want in ((hc, ref_hc), (normed, ref_normed)): + assert got.dtype == want.dtype and got.shape == want.shape + if got.numel(): + ratio = _ulp_ratio(got, want) + assert ratio <= 1.0, f"{ratio:.4g} ulp" + + +def _combine(block_out, shared, gate_logits) -> torch.Tensor: + """The MoE block output in float64, rounded once to `block_out`'s dtype.""" + s = torch.sigmoid(gate_logits.double().reshape(-1, 1)) + return (block_out.double() + s * shared.double()).to(block_out.dtype) + + +def _routed(block_out, topk_weights, topk_rows): + """The routed rows: `block_out` itself, or with the fold the float64 finalize of + the permuted rows rounded once.""" + if topk_rows is None: + return block_out + valid = (topk_rows >= 0).unsqueeze(-1) + picked = block_out.double()[topk_rows.clamp(min=0).long()] + w = topk_weights.double().unsqueeze(-1) + return torch.where(valid, picked * w, 0.0).sum(1).to(block_out.dtype) + + +def _compare_moe(args, out, ref) -> None: + """MoE-combine cells. + + `hc_out`: the op's combine is an fp32 value rounded once, the reference's a + float64 value rounded once; where the two straddle a rounding midpoint they + differ by one ulp of the block output, and the accumulate carries that + scaled by `inject`. So the band is one ulp of `hc_out` plus one ulp of the + combine times `|inject|`, element by element -- still well under one ulp + of `hc_out` wherever the block term is small against the residual. With the + finalize fold the routed row has its own rounding (fp32 FMAs against the + float64 sum), one more ulp of the routed row times `|inject|`. + + `normed`: one ulp of the float64 norm of the residual the op stored, + `reference(hc_out, None, None, ...)` -- the dataflow the op states.""" + hyper_input, block_out, inject, gamma, eps, c, shared, gate = args[:8] + topk_weights, topk_rows = args[8:] if len(args) > 8 else (None, None) + (hc, normed), (ref_hc, _) = out, ref + assert hc.dtype == ref_hc.dtype and hc.shape == ref_hc.shape + if not hc.numel(): + return + feps, tiny = torch.finfo(ref_hc.dtype).eps, torch.finfo(ref_hc.dtype).tiny + routed = _routed(block_out, topk_weights, topk_rows) + b = _combine(routed, shared, gate).double().abs() + if topk_rows is not None: + b = b + routed.double().abs() + band = feps * ref_hc.double().abs() + ( + feps * b.unsqueeze(-2) * inject.double().abs().unsqueeze(-1) + ).flatten(-2) + ratio = ((hc.double() - ref_hc.double()).abs() / (band + tiny)).max().item() + assert ratio <= 1.0, f"hc_out {ratio:.4g}x the band" + own = op.reference(hc, None, None, gamma, eps, c)[1] + assert _ulp_ratio(normed, own) <= 1.0, f"normed {_ulp_ratio(normed, own):.4g} ulp" + + +def _build(spec, seed): + rows = spec["rows"] + hc_count = spec.get("hc_count", HC_COUNT) + hidden = spec.get("hidden", HIDDEN) + dtype = spec.get("dtype", torch.bfloat16) + width = hc_count * hidden + g = torch.Generator(device=DEV).manual_seed(seed) + hyper_input = (torch.randn(rows, width, generator=g, device=DEV) * 3.0).to(dtype) + if spec.get("layout") == "row_slice": + buf = torch.randn(2 * rows, width, generator=g, device=DEV).to(dtype) + buf[::2] = hyper_input + hyper_input = buf[::2] + gamma = torch.randn(width, generator=g, device=DEV) * GAMMA_RMS + gamma[0], gamma[width // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX + gamma = gamma.to(dtype) + block_out = inject = shared = gate_logits = None + if spec.get("inject", True): + block_out = (torch.randn(rows, hidden, generator=g, device=DEV) * 2.0).to(dtype) + inject = (2.0 * torch.sigmoid(torch.randn(rows, hc_count, generator=g, device=DEV))).to( + dtype + ) + if spec.get("moe", False): + shared = (torch.randn(rows, hidden, generator=g, device=DEV) * 2.0).to(dtype) + gate_logits = (torch.randn(rows, 1, generator=g, device=DEV) * 2.0).to(dtype) + if spec.get("gate") == "flat": + gate_logits = gate_logits.reshape(rows) + top_k = spec.get("fold", 0) + if not top_k: + return hyper_input, block_out, inject, gamma, EPS, hc_count, shared, gate_logits + # The runner's unfinalized output: permuted FC2 rows (more rows than slots, as + # its padded buffer has), each slot's row and its routing weight. + n_perm = rows * top_k + 37 + block_out = (torch.randn(n_perm, hidden, generator=g, device=DEV) * 2.0).to(dtype) + order = torch.randperm(n_perm, generator=g, device=DEV)[: rows * top_k] + topk_rows = order.reshape(rows, top_k).to(torch.int32) + if spec.get("drop"): + topk_rows[:, -3:] = -1 + topk_weights = torch.softmax(torch.randn(rows, top_k, generator=g, device=DEV), -1).to(dtype) + return ( + hyper_input, + block_out, + inject, + gamma, + EPS, + hc_count, + shared, + gate_logits, + topk_weights, + topk_rows, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) + before = [a.clone() if isinstance(a, torch.Tensor) else a for a in args] + hc, normed = op(*args) + for a, b in zip(args, before): + assert not isinstance(a, torch.Tensor) or torch.equal(a, b), "every input is read-only" + if args[1] is None: + assert hc is args[0], "norm-only returns hyper_input itself" + else: + assert hc.is_contiguous() + assert args[0].numel() == 0 or hc.data_ptr() != args[0].data_ptr(), "hc_out is fresh" + assert normed.is_contiguous() + if args[6] is None: + _compare((hc, normed), op.reference(*args)) + else: + _compare_moe(args, (hc, normed), op.reference(*args)) + + +def _old_composition(hyper_input, block_out, inject, gamma, eps, hc_count): + hc = qwen4_exp_post_inject(hyper_input, block_out, inject) + hidden = hc.shape[1] // hc_count + parts = torch.split(hc, [hidden] * hc_count, 1) + weights = torch.split(gamma, hidden, 0) + normed = torch.cat([flashinfer_gemma_rmsnorm(p, w, eps) for w, p in zip(weights, parts)], 1) + return hc, normed + + +@pytest.mark.parametrize("rows", [1, 128, 8192]) +def test_matches_the_composition_it_replaces(rows) -> None: + """Against post_inject + four gemma norms + cat: the residual within one + ulp (fma contraction only), the normed stream within one ulp.""" + args = _build(dict(rows=rows), seed=rows)[:6] + hc, normed = op(*args) + old_hc, old_normed = _old_composition(*args) + assert _ulp_ratio(hc, old_hc) <= 1.0 + assert _ulp_ratio(normed, old_normed) <= 1.0 + print(f"\n rows {rows}: hc bitwise-equal fraction {(hc == old_hc).float().mean().item():.6f}") + + +def test_wrong_variants_land_outside_the_band() -> None: + """Each mis-reading keeps the shapes; each lands > 5x outside one ulp. The + streams carry different scales, as the trained residual's do, so a norm + over the wrong group is visible.""" + hyper_input, block_out, inject, gamma, eps, c, _, _ = _build(dict(rows=64), seed=5) + scales = torch.tensor([0.5, 1.0, 2.0, 4.0], device=DEV).repeat_interleave(HIDDEN) + hyper_input = (hyper_input.float() * scales).to(torch.bfloat16) + want_hc, want_normed = op.reference(hyper_input, block_out, inject, gamma, eps, c) + x = hyper_input.double().unflatten(-1, (c, HIDDEN)) + b, g = block_out.double(), inject.double() + wrong_hc = { + "stream-minor inject": (x.flatten(-2) + (b.unsqueeze(-1) * g.unsqueeze(-2)).flatten(-2)), + "no gate": (x + b.unsqueeze(-2)).flatten(-2), + } + for name, value in wrong_hc.items(): + ratio = _ulp_ratio(value.to(torch.bfloat16), want_hc) + print(f"\n hc {name:<24} {ratio:10.4g} ulp") + assert ratio > 5.0, name + hc = want_hc.double().unflatten(-1, (c, HIDDEN)) + full = hc.flatten(-2) + wrong_normed = { + "RMS over the full width": full + * torch.rsqrt(full.pow(2).mean(-1, keepdim=True) + eps) + * (1.0 + gamma.double()), + "w instead of 1 + w": (hc * torch.rsqrt(hc.pow(2).mean(-1, keepdim=True) + eps)).flatten(-2) + * gamma.double(), + "norm of the pre-inject stream": (x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)) + .flatten(-2) + .mul(1.0 + gamma.double()), + } + for name, value in wrong_normed.items(): + ratio = _ulp_ratio(value.to(torch.bfloat16), want_normed) + print(f" normed {name:<30} {ratio:10.4g} ulp") + assert ratio > 5.0, name + + +def test_guard_refuses_block_out_without_inject() -> None: + hyper_input, block_out, inject, gamma, eps, c, _, _ = _build(dict(rows=4), seed=6) + with validating(op), pytest.raises(AssertionError): + op(hyper_input, block_out, None, gamma, eps, c) + with validating(op), pytest.raises(AssertionError): + op(hyper_input, block_out, inject[0], gamma, eps, c) + + +def test_cuda_graph_capture_and_replay() -> None: + args = list(_build(dict(rows=8), seed=8)) + static_in = args[0] + 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): + graph_hc, graph_normed = op(*args) + fresh = _build(dict(rows=8), seed=18)[0] + static_in.copy_(fresh) + graph.replay() + torch.cuda.synchronize() + eager_hc, eager_normed = op(fresh, *args[1:]) + assert torch.equal(graph_hc, eager_hc) + assert torch.equal(graph_normed, eager_normed) + + +@pytest.mark.parametrize("rows", [2, 128, 2048]) +def test_fold_matches_the_runner_finalize(rows) -> None: + """With the fold, bitwise this op on the runner's finalized rows: `finalize` + (the permuted MoE entry's emulation of the runner's finalize, bitwise its + output) and then the combine without the fold.""" + args = _build(dict(rows=rows, moe=True, fold=10), seed=300 + rows) + hyper_input, rows_, inject, gamma, eps, c, shared, gate, topk_weights, topk_rows = args + hc, normed = op(*args) + routed = finalize(rows_, topk_weights, topk_rows) + old_hc, old_normed = op(hyper_input, routed, inject, gamma, eps, c, shared, gate) + assert torch.equal(hc, old_hc) + assert torch.equal(normed, old_normed) + + +@pytest.mark.parametrize("rows", [1, 128, 8192]) +def test_moe_combine_matches_the_composition_it_replaces(rows) -> None: + """Against `fused_sigmoid_gate_mul_add` (in place on the routed output, as + the forward called it) and then this op without the combine: bitwise.""" + args = _build(dict(rows=rows, moe=True), seed=100 + rows) + hyper_input, block_out, inject, gamma, eps, c, shared, gate = args + hc, normed = op(*args) + combined = fused_sigmoid_gate_mul_add(block_out.clone(), gate, shared) + old_hc, old_normed = op(hyper_input, combined, inject, gamma, eps, c) + print( + f"\n rows {rows}: hc bitwise {(hc == old_hc).float().mean().item():.6f}, " + f"normed bitwise {(normed == old_normed).float().mean().item():.6f}" + ) + assert torch.equal(hc, old_hc) + assert torch.equal(normed, old_normed) + + +def test_moe_wrong_variants_land_outside_the_band() -> None: + """Each mis-reading of the combine keeps the shapes; each lands > 5x outside + the `_compare_moe` residual band.""" + args = _build(dict(rows=64, moe=True), seed=7) + hyper_input, block_out, inject, gamma, eps, c, shared, gate = args + want_hc = op.reference(*args)[0] + feps = torch.finfo(torch.bfloat16).eps + b = _combine(block_out, shared, gate).double().abs() + band = feps * want_hc.double().abs() + ( + feps * b.unsqueeze(-2) * inject.double().abs().unsqueeze(-1) + ).flatten(-2) + r, sh, s = block_out.double(), shared.double(), torch.sigmoid(gate.double()) + variants = { + "no sigmoid": r + gate.double() * sh, + "gate on the routed half": s * r + sh, + "gate dropped": r + sh, + "shared dropped": r, + "1 - sigmoid": r + (1.0 - s) * sh, + } + x = hyper_input.double().unflatten(-1, (c, HIDDEN)) + for name, blk in variants.items(): + wrong = (x + blk.unsqueeze(-2) * inject.double().unsqueeze(-1)).flatten(-2) + ratio = ((wrong.to(torch.bfloat16).double() - want_hc.double()).abs() / band).max().item() + print(f"\n hc {name:<26} {ratio:10.4g}x the band") + assert ratio > 5.0, name + + +def test_guard_refuses_a_partial_combine() -> None: + hyper_input, block_out, inject, gamma, eps, c, shared, gate = _build( + dict(rows=4, moe=True), seed=9 + ) + fold = _build(dict(rows=4, moe=True, fold=10), seed=9) + with validating(op), pytest.raises(AssertionError): + op(*fold[:6], None, None, *fold[8:]) + with validating(op), pytest.raises(AssertionError): + op(*fold[:9], None) + with validating(op), pytest.raises(AssertionError): + op(*fold[:9], fold[9].long()) + with validating(op), pytest.raises(AssertionError): + op(hyper_input, block_out, inject, gamma, eps, c, shared, None) + with validating(op), pytest.raises(AssertionError): + op(hyper_input, None, None, gamma, eps, c, shared, gate) + with validating(op), pytest.raises(AssertionError): + op(hyper_input, block_out, inject, gamma, eps, c, shared, gate[:1].expand(4, 1)) + + +def test_moe_cuda_graph_capture_and_replay() -> None: + args = list(_build(dict(rows=8, moe=True), seed=10)) + 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): + graph_hc, graph_normed = op(*args) + fresh = _build(dict(rows=8, moe=True), seed=20) + for i in (0, 1, 6, 7): + args[i].copy_(fresh[i]) + graph.replay() + torch.cuda.synchronize() + eager_hc, eager_normed = op(*args) + assert torch.equal(graph_hc, eager_hc) + assert torch.equal(graph_normed, eager_normed) + + +@pytest.mark.parametrize( + "spec", + [ + dict(rows=1), + dict(rows=2, drop=True), + dict(rows=8), + dict(rows=33, drop=True), + dict(rows=128), + dict(rows=17, gate="flat", layout="row_slice"), + dict(rows=300, drop=True), + dict(rows=511), + ], + ids=lambda s: "-".join(f"{k}{v}" for k, v in s.items()), +) +def test_fold_cute_is_the_triton_kernel_bitwise(spec, monkeypatch) -> None: + """Below `_FEW_WARPS_MIN_ROWS` rows the fold runs the CuTe DSL kernel; its + `hc_out` and `normed` are bitwise the Triton kernel's (run by refusing the + CuTe guard).""" + args = _build(dict(spec, moe=True, fold=10), seed=500 + spec["rows"]) + calls = [] + cute_fold = M._fold_cute + + def recording(*a): + calls.append(a[0].shape[0]) + cute_fold(*a) + + monkeypatch.setattr(M, "_fold_cute", recording) + hc, normed = op(*args) + assert calls == [spec["rows"]], "the decode MoE site runs the CuTe kernel" + monkeypatch.setattr(M, "_fold_cute_ok", lambda *a: False) + want_hc, want_normed = op(*args) + assert torch.equal(hc, want_hc) + assert torch.equal(normed, want_normed) + + +def test_fold_from_the_few_warps_rows_keeps_triton(monkeypatch) -> None: + """From `_FEW_WARPS_MIN_ROWS` rows (the Triton kernel's two-warp form, another + reduction order) the fold stays on the Triton kernel.""" + args = _build(dict(rows=M._FEW_WARPS_MIN_ROWS, moe=True, fold=10), seed=77) + monkeypatch.setattr(M, "_fold_cute", lambda *a: pytest.fail("CuTe kernel at 512 rows")) + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_lowrank_act.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_lowrank_act.py new file mode 100644 index 000000000000..8226d441001d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_lowrank_act.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qwen4_exp_hc_lowrank_act entry: cells within one ulp of the +float64 reference, wrong formulas outside, graph capture.""" + +import zlib + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_lowrank_act import ( # noqa: E501 + qwen4_exp_hc_lowrank_act as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "qwen4_exp_hc_lowrank_act requires a CUDA device" + +DEV = "cuda" +RANK = 320 + + +def _ulp_ratio(out, ref) -> float: + eps = torch.finfo(ref.dtype).eps + tiny = torch.finfo(ref.dtype).tiny + return ((out.double() - ref.double()).abs() / (eps * ref.double().abs() + tiny)).max().item() + + +def _compare(out, ref) -> None: + """One ulp: each output is one fp32 activation of an fp32 input, rounded + once, so it can differ from the once-rounded float64 value only at a + rounding midpoint.""" + for got, want in zip(out, ref): + assert (got is None) == (want is None) + if got is None: + continue + assert got.dtype == want.dtype and got.shape == want.shape and got.is_contiguous() + if got.numel(): + assert _ulp_ratio(got, want) <= 1.0, f"{_ulp_ratio(got, want):.4g} ulp" + + +def _build(spec, seed): + hc_count = spec.get("hc_count", 4) + g = torch.Generator(device=DEV).manual_seed(seed) + # The first projection's logits on this checkpoint reach |67| (down) / |48| (inject). + logits = torch.randn(spec["rows"], spec.get("logits_width", 384), generator=g, device=DEV) * 20 + return logits, RANK, hc_count, spec.get("inject", True), spec.get("dtype", torch.bfloat16) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) + low, inject = op(*args) + assert low.shape == (cell.spec["rows"], RANK) + _compare((low, inject), op.reference(*args)) + + +def test_wrong_variants_land_outside_the_band() -> None: + logits, rank, c, wi, dt = _build(dict(rows=64), seed=3) + want_low, want_inject = op.reference(logits, rank, c, wi, dt) + lg = logits.double() + for name, value, want in ( + ("silu without /hc_count", F.silu(lg[:, :rank]), want_low), + ("sigmoid instead of silu", torch.sigmoid(lg[:, :rank] / c), want_low), + ("inject missing the 2", torch.sigmoid(lg[:, rank : rank + c] / c), want_inject), + ( + "inject from the padding", + 2 * torch.sigmoid(lg[:, rank + c : rank + 2 * c] / c), + want_inject, + ), + ): + ratio = _ulp_ratio(value.to(dt), want) + print(f"\n {name:<28} {ratio:10.4g} ulp") + assert ratio > 5.0, name + + +def test_guard_refuses_bf16_logits() -> None: + logits, rank, c, wi, dt = _build(dict(rows=4), seed=4) + with validating(op), pytest.raises(AssertionError): + op(logits.to(torch.bfloat16), rank, c, wi, dt) + + +def test_cuda_graph_capture_and_replay() -> None: + logits, rank, c, wi, dt = _build(dict(rows=8), seed=5) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + op(logits, rank, c, wi, dt) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + g_low, g_inject = op(logits, rank, c, wi, dt) + logits.copy_(_build(dict(rows=8), seed=6)[0]) + graph.replay() + torch.cuda.synchronize() + e_low, e_inject = op(logits, rank, c, wi, dt) + assert torch.equal(g_low, e_low) and torch.equal(g_inject, e_inject) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py new file mode 100644 index 000000000000..9fa943fee089 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py @@ -0,0 +1,150 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qwen4_exp_hc_up_gate_mean entry: cells in the band +`qwen4_exp_gated_residual`'s test derives for `mixed`, wrong formulas outside, +the guard, graph capture, and the mixer as the forward runs it against the +gated-residual float64 reference.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501 + gate_ratio, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_up_gate_mean import ( # noqa: E501 + qwen4_exp_hc_up_gate_mean as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "qwen4_exp_hc_up_gate_mean requires a CUDA device" + +DEV = "cuda" + + +def _compare(out, ref) -> None: + """`rtol = eps`, `atol = eps * max|ref|` (`gate_ratio <= 1`), the band + `qwen4_exp_gated_residual` derives for `mixed`.""" + assert out.dtype == ref.dtype and out.shape == ref.shape and out.is_contiguous() + if out.numel(): + assert gate_ratio(out, ref) <= 1.0, f"{gate_ratio(out, ref):.4g}x the band" + + +def _build(spec, seed): + hc, hidden, rank = spec.get("hc_count", 4), spec.get("hidden", 2560), spec.get("rank", 320) + g = torch.Generator(device=DEV).manual_seed(seed) + rows, width = spec["rows"], hc * hidden + # This checkpoint's statistics: low = silu(.) of the bottleneck (|x| < ~3), the up + # weight ~N(0, 0.068) (gate logits within about +-8), the normed stream ~+-30. + low = torch.nn.functional.silu(torch.randn(rows, rank, generator=g, device=DEV) * 1.5).to( + torch.bfloat16 + ) + up_w = (torch.randn(width, rank, generator=g, device=DEV) * 0.068).to(torch.bfloat16) + normed = (torch.randn(rows, width, generator=g, device=DEV) * 6).to(torch.bfloat16) + return low, up_w.t(), normed, hc + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) + mixed = op(*args) + assert mixed.shape == (cell.spec["rows"], cell.spec.get("hidden", 2560)) + _compare(mixed, op.reference(*args)) + + +def test_wrong_variants_land_outside_the_band() -> None: + low, up, normed, c = _build(dict(rows=128), seed=3) + want = op.reference(low, up, normed, c) + got = op(low, up, normed, c) + h = normed.shape[1] // c + logits = (low.double() @ up.double()).unflatten(-1, (c, h)) + x = normed.double().unflatten(-1, (c, h)) + g = torch.sigmoid(logits) + for name, value in ( + ("sum instead of mean", (g * x).sum(-2)), + ("no sigmoid", (logits * x).mean(-2)), + ("stream-minor packing", (g.transpose(-1, -2).reshape(-1, c, h) * x).mean(-2)), + ("ungated mean", x.mean(-2)), + ): + ratio = gate_ratio(value.to(want.dtype), want) + print(f"\n {name:<24} {ratio:10.4g}x the band (op: {gate_ratio(got, want):.3g}x)") + assert ratio > 5.0, name + + +def test_guard_refuses_a_contiguous_up() -> None: + low, up, normed, c = _build(dict(rows=4), seed=4) + with validating(op), pytest.raises(AssertionError): + op(low, up.contiguous(), normed, c) + + +def test_cuda_graph_capture_and_replay() -> None: + low, up, normed, c = _build(dict(rows=8), seed=5) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + op(low, up, normed, c) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + g_mixed = op(low, up, normed, c) + fresh = _build(dict(rows=8), seed=6) + low.copy_(fresh[0]) + normed.copy_(fresh[2]) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(g_mixed, op(fresh[0], up, fresh[2], c)) + + +@pytest.mark.parametrize("rows", [1, 33, 128, 2048, 8192, 16384]) +@pytest.mark.parametrize("with_inject", [True, False]) +def test_mixer_matches_the_gated_residual_reference(rows, with_inject) -> None: + """The mixer as the forward runs it -- `cublas_mm` over the zero-padded + `[down; inject]` (fp32 out), `qwen4_exp_hc_lowrank_act`, this op -- against + `qwen4_exp_gated_residual.reference` (float64 from `normed`), in the band + that entry certifies its own op in.""" + from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm + from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501 + qwen4_exp_gated_residual, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_gate_mean import ( # noqa: E501 + qwen4_exp_hc_gate_mean, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_lowrank_act import ( # noqa: E501 + qwen4_exp_hc_lowrank_act, + ) + + hc, hidden, rank, eps = 4, 2560, 320, 1e-6 + width = hc * hidden + g = torch.Generator(device=DEV).manual_seed(rows) + + def normal(*shape, scale): + return (torch.randn(*shape, generator=g, device=DEV) * scale).to(torch.bfloat16) + + # This checkpoint's layer-0 hyper-connection element statistics. + down, up = normal(rank, width, scale=0.0296), normal(width, rank, scale=0.0681) + inject_w, gamma = normal(hc, width, scale=0.0107), normal(width, scale=0.477) + gamma[0], gamma[width // 2] = 6.625, -6.625 + x = torch.randn(rows, hc, hidden, generator=g, device=DEV) + x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) + normed = (x.flatten(-2) * (1.0 + gamma.float())).to(torch.bfloat16) + first = torch.cat([down, inject_w], 0) if with_inject else down + first = torch.cat( + [first, torch.zeros(384 - first.shape[0], width, device=DEV, dtype=first.dtype)] + ) + low, _ = qwen4_exp_hc_lowrank_act( + cublas_mm(normed, first.t(), None, torch.float32), rank, hc, with_inject, torch.bfloat16 + ) + mixed = op(low, up.t(), normed, hc) + want = qwen4_exp_gated_residual.reference( + normed, down, up, inject_w if with_inject else None, hc + ) + _compare(mixed, want[0]) + prev = qwen4_exp_hc_gate_mean(cublas_mm(low, up.t()), normed, hc) + print( + f"\n rows {rows}: fused {gate_ratio(mixed, want[0]):.3f}x band, " + f"cublas + gate_mean {gate_ratio(prev, want[0]):.3f}x" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4.py new file mode 100644 index 000000000000..a315944a3393 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4.py @@ -0,0 +1,183 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qwen4_exp_hc_up_gate_mean_fp4 entry: `mixed` in the band +`qwen4_exp_hc_up_gate_mean`'s test uses and bitwise that op's, the NVFP4 +outputs bitwise against `torch.ops.trtllm.fp4_quantize` of the returned +`mixed` and against the native-torch quantization outside the reciprocal tie +window, wrong recipes outside, guards, graph capture.""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501 + gate_ratio, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_decode_mix_fp4 import ( # noqa: E501 + nvfp4_near_tie, + nvfp4_reference, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_up_gate_mean import ( # noqa: E501 + qwen4_exp_hc_up_gate_mean, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_up_gate_mean_fp4 import ( # noqa: E501 + qwen4_exp_hc_up_gate_mean_fp4 as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "qwen4_exp_hc_up_gate_mean_fp4 requires a CUDA device" + +DEV = "cuda" + + +def _unpack(xq: torch.Tensor) -> torch.Tensor: + return torch.stack([xq & 0xF, xq >> 4], -1).flatten(-2) + + +def _check_quant(mixed, xq, xsf, gs) -> None: + """NVFP4 outputs against the native-torch quantization of the `mixed` this + call returned: scale bytes equal, e2m1 codes equal except inside the + reciprocal tie window.""" + want_q, want_sf = nvfp4_reference(mixed, gs) + assert xq.dtype == torch.uint8 and xq.shape == want_q.shape and xq.is_contiguous() + assert xsf.dtype == torch.uint8 and xsf.shape == want_sf.shape + assert torch.equal(xsf, want_sf), f"{(xsf != want_sf).sum().item()} scale bytes differ" + diff = _unpack(xq) != _unpack(want_q) + stray = diff & ~nvfp4_near_tie(mixed, gs) + assert not stray.any(), f"{stray.sum().item()} e2m1 codes differ outside the tie window" + + +def _compare(out, ref, gs) -> None: + """`mixed`: gate_ratio <= 1, the band `qwen4_exp_hc_up_gate_mean` is + certified in. Quantization: exact on the returned `mixed`.""" + (mixed, xq, xsf), (ref_mixed, _, _) = out, ref + assert mixed.dtype == ref_mixed.dtype and mixed.shape == ref_mixed.shape + assert mixed.is_contiguous() + if mixed.numel() == 0: + assert xq.numel() == 0 and xsf.numel() == 0 + return + assert gate_ratio(mixed, ref_mixed) <= 1.0, f"{gate_ratio(mixed, ref_mixed):.4g}x the band" + _check_quant(mixed, xq, xsf, gs) + + +def _build(spec, seed): + hc, hidden, rank = spec.get("hc_count", 4), spec.get("hidden", 2560), spec.get("rank", 320) + g = torch.Generator(device=DEV).manual_seed(seed) + rows, width = spec["rows"], hc * hidden + # This checkpoint's statistics: low = silu(.) of the bottleneck (|x| < ~3), the up + # weight ~N(0, 0.068) (gate logits within about +-8), the normed stream ~+-30. + low = torch.nn.functional.silu(torch.randn(rows, rank, generator=g, device=DEV) * 1.5).to( + torch.bfloat16 + ) + up_w = (torch.randn(width, rank, generator=g, device=DEV) * 0.068).to(torch.bfloat16) + normed = (torch.randn(rows, width, generator=g, device=DEV) * 6).to(torch.bfloat16) + # A calibrated activation global scale: 448 * 6 / amax over a typical batch. + gs = torch.tensor([448.0 * 6.0 / 24.0], device=DEV, dtype=torch.float32) + return low, up_w.t(), normed, hc, gs + + +def _seed(cell) -> int: + return zlib.crc32(cell.why.encode()) % (2**31) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + certifying(op) + args = _build(cell.spec, seed=_seed(cell)) + out = op(*args) + assert out[0].shape == (cell.spec["rows"], cell.spec.get("hidden", 2560)) + _compare(out, op.reference(*args), args[4]) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_bitwise_the_composition_it_replaces(cell) -> None: + """`qwen4_exp_hc_up_gate_mean` then `torch.ops.trtllm.fp4_quantize` (linear + scales): every output bitwise equal, ties included.""" + low, up, normed, c, gs = _build(cell.spec, seed=_seed(cell)) + mixed, xq, xsf = op(low, up, normed, c, gs) + old = qwen4_exp_hc_up_gate_mean(low, up, normed, c) + assert torch.equal(mixed, old) + if old.shape[0] == 0: + return + old_q, old_sf = torch.ops.trtllm.fp4_quantize(old, gs, 16, False, False) + assert torch.equal(xq, old_q.view(torch.uint8)), ( + f"{(xq != old_q.view(torch.uint8)).sum().item()} bytes" + ) + assert torch.equal(xsf, old_sf.view(torch.uint8).reshape(-1)) + + +@pytest.mark.parametrize("scale", [1e-6, 1e-3, 1.0, 448.0 * 6.0 / 24.0, 1e4]) +def test_bitwise_across_global_scales(scale) -> None: + """Tiny scales (zero / subnormal e4m3 bytes) and saturating ones.""" + low, up, normed, c, _ = _build(dict(rows=128), seed=11) + h = normed.shape[1] // c + for s in range(c): # row 0's first 64 mix columns are exactly zero: amax 0 + normed[0, s * h : s * h + 64] = 0 + gs = torch.tensor([scale], device=DEV, dtype=torch.float32) + mixed, xq, xsf = op(low, up, normed, c, gs) + assert not mixed[0, :64].any() + old_q, old_sf = torch.ops.trtllm.fp4_quantize(mixed, gs, 16, False, False) + assert torch.equal(xq, old_q.view(torch.uint8)) + assert torch.equal(xsf, old_sf.view(torch.uint8).reshape(-1)) + + +def test_wrong_recipes_fail_the_quant_check() -> None: + """Each plausible wrong layout or recipe is refused by `_check_quant`.""" + low, up, normed, c, gs = _build(dict(rows=128), seed=3) + mixed, xq, xsf = op(low, up, normed, c, gs) + _check_quant(mixed, xq, xsf, gs) + rows = mixed.shape[0] + swapped = ((xq & 0xF) << 4) | (xq >> 4) + sf2 = xsf.view(rows, -1) + # 128 rows: the 128x4-swizzled buffer has exactly the linear one's length. + sw = torch.ops.trtllm.fp4_quantize(mixed, gs, 16, False, True)[1].view(torch.uint8).reshape(-1) + h = normed.shape[1] // c + unrounded = ( + torch.sigmoid(low.float() @ up.float()).unflatten(-1, (c, h)) + * normed.float().unflatten(-1, (c, h)) + ).mean(-2) + q32, sf32 = nvfp4_reference(unrounded, gs) + for name, args in ( + ("nibbles swapped", (mixed, swapped, xsf)), + ("swizzled scales", (mixed, xq, sw)), + ("scale one code high", (mixed, xq, (sf2 + 1).reshape(-1))), + ("quantized before the bf16 rounding", (mixed, q32, sf32)), + ): + with pytest.raises(AssertionError): + _check_quant(*args, gs) + print(f"\n {name:<36} refused") + + +def test_guards() -> None: + low, up, normed, c, gs = _build(dict(rows=4), seed=4) + with validating(op), pytest.raises(AssertionError): + op(low, up.contiguous(), normed, c, gs) + with validating(op), pytest.raises(AssertionError): + op(low, up, normed, c, torch.cat([gs, gs])) + with validating(op), pytest.raises(AssertionError): + op(low, up, normed, c, gs.double()) + + +def test_cuda_graph_capture_and_replay() -> None: + low, up, normed, c, gs = _build(dict(rows=8), seed=5) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + op(low, up, normed, c, gs) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + g_out = op(low, up, normed, c, gs) + fresh = _build(dict(rows=8), seed=6) + low.copy_(fresh[0]) + normed.copy_(fresh[2]) + gs.mul_(0.5) + graph.replay() + torch.cuda.synchronize() + want = op(fresh[0], up, fresh[2], c, gs) + for got, exp in zip(g_out, want): + assert torch.equal(got, exp) From a14b0b995493dfa38a2f001e7f1e844c21b64823 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Sun, 18 Oct 2026 23:00:04 -0700 Subject: [PATCH 08/28] [TRTLLM-16460][perf] Qwen3.8-Flash-Next target: MoE, GEMM, PLE and elementwise ops The target's remaining fast-path ops, with their unit tests: - MoE: fp4_block_scale_moe_renorm_routed / _permuted (the trtllm-gen runner entered with router logits, finalized or left permuted), fp4_moe_renorm_cute_dsl (prefill-sized launches on TRT-LLM's CuTe DSL MoE kernels); updates to fp8_block_scale_moe_runner and renorm_moe_routing_op. - bf16_gemv_pdl: decode-row bf16 projection under PDL, its weight streaming in while the kernel before it finishes. - PLE: ple_ngram_embed (id history, n-gram hash, table gather and fp8 dequantize in one launch) and ple_gate_conv_inject_norm (the PLE block tail for decode-shaped batches). - silu_and_mul_rows, bf16_to_fp32 (the logits widening), updates to fused_sigmoid_mul_inplace and fp8_quantize_1x128. Folds these commits of the review branch (their changes to the files of this commit): - qwen3.8 target: carry the decode-fusion perf work onto the feature target (62151b727b) - qwen3.8 target: shared-expert gate logit from the gate|up GEMM at prefill (2ad541143d) - qwen3.8 target: widen the logits to fp32 in one vectorized kernel (edb0d2e93a) Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../activation/fused_sigmoid_mul_inplace.py | 12 + .../ops/activation/silu_and_mul_rows.py | 92 +++ .../ops/gemm/__init__.py | 2 + .../ops/gemm/bf16_gemv_pdl.py | 392 ++++++++++ .../fp4_block_scale_moe_renorm_permuted.py | 144 ++++ .../moe/fp4_block_scale_moe_renorm_routed.py | 225 ++++++ .../ops/moe/fp4_moe_renorm_cute_dsl.py | 220 ++++++ .../ops/moe/fp8_block_scale_moe_runner.py | 12 +- .../ops/moe/renorm_moe_routing_op.py | 3 +- .../ops/ple/ple_gate_conv_inject_norm.py | 676 ++++++++++++++++++ .../ops/ple/ple_ngram_embed.py | 591 +++++++++++++++ .../ops/quantization/bf16_to_fp32.py | 84 +++ .../ops/quantization/fp8_quantize_1x128.py | 9 +- ...t_modeling_v2_fused_sigmoid_mul_inplace.py | 5 + .../test_modeling_v2_silu_and_mul_rows.py | 75 ++ .../gemm/test_modeling_v2_bf16_gemv_pdl.py | 129 ++++ ..._v2_fp4_block_scale_moe_renorm_permuted.py | 118 +++ ...ng_v2_fp4_block_scale_moe_renorm_routed.py | 152 ++++ ...t_modeling_v2_ple_gate_conv_inject_norm.py | 298 ++++++++ .../test_modeling_v2_bf16_to_fp32.py | 69 ++ 20 files changed, 3303 insertions(+), 5 deletions(-) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/silu_and_mul_rows.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/bf16_gemv_pdl.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_permuted.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_routed.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_moe_renorm_cute_dsl.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_gate_conv_inject_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_ngram_embed.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/bf16_to_fp32.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_silu_and_mul_rows.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_bf16_gemv_pdl.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_block_scale_moe_renorm_permuted.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_gate_conv_inject_norm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_bf16_to_fp32.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/fused_sigmoid_mul_inplace.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/fused_sigmoid_mul_inplace.py index 4982ae3f00fe..b870dd1627c1 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/fused_sigmoid_mul_inplace.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/fused_sigmoid_mul_inplace.py @@ -45,6 +45,18 @@ class FusedSigmoidMulInplace(OpWrapper): why="prefill 2048 tokens: >= 1024 rows selects 2048-element blocks", spec=dict(tokens=2048, **_CKPT), ), + Cell( + why="decode 1 token: 2-D gate, the gate columns of the fused projection row", + spec=dict(tokens=1, heads=24, head_dim=256, dtype=torch.bfloat16, gate="fused_row"), + ), + Cell( + why="decode 128 tokens: 2-D gate off the fused projection row (stride 14080)", + spec=dict(tokens=128, heads=24, head_dim=256, dtype=torch.bfloat16, gate="fused_row"), + ), + Cell( + why="prefill 8192 tokens: 2-D gate off the fused projection row", + spec=dict(tokens=8192, heads=24, head_dim=256, dtype=torch.bfloat16, gate="fused_row"), + ), Cell( why="attention_output a [33, 6144] window of a [35, 8192] buffer", spec=dict(tokens=33, out="window", **_CKPT), diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/silu_and_mul_rows.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/silu_and_mul_rows.py new file mode 100644 index 000000000000..ad0d9457514d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/silu_and_mul_rows.py @@ -0,0 +1,92 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""SiLU-and-mul over the leading `2 * inter` columns of each row, read through the row stride: +the shared expert's activation taken out of its fused gate|up|gate-logit GEMM output. + +Tier 5 (Triton written for this target). `flashinfer_silu_and_mul` addresses row `t` at +`t * 2 * inter`, so its input must be exactly the gate|up columns, contiguous; here they are +the head of a wider row (the same GEMM also forms the per-token gate logit) and are read in +place. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: This checkpoint's shared expert: intermediate 640, the fused GEMM output row 2 * 640 + 8. +_INTER = 640 +_WIDTH = 2 * _INTER + 8 +_BLOCK = 1024 + + +@triton.jit +def _silu_and_mul_rows_kernel(x_ptr, out_ptr, stride_x, INTER: tl.constexpr, BLOCK: tl.constexpr): + row = tl.program_id(0).to(tl.int64) + cols = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + mask = cols < INTER + g = tl.load(x_ptr + row * stride_x + cols, mask=mask, other=0.0).to(tl.float32) + u = tl.load(x_ptr + row * stride_x + INTER + cols, mask=mask, other=0.0).to(tl.float32) + y = g / (1.0 + tl.exp(-g)) * u + tl.store(out_ptr + row * INTER + cols, y.to(out_ptr.dtype.element_ty), mask=mask) + + +class SiluAndMulRows(OpWrapper): + """`out[t, i] = silu(x[t, i]) * x[t, inter + i]` for `i < inter`: the gate half first, as + `flashinfer_silu_and_mul` reads it, in fp32 and rounded once to `x`'s dtype. Columns of + `x` past `2 * inter` are not read. Returns a fresh contiguous `[T, inter]` tensor. + + Fusion boundary: the activation only; the GEMM that forms `x` and the one that reads + `out` are the caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="prefill chunk 16384 rows of the fused 1288-wide output", spec=dict(rows=16384)), + Cell(why="decode batch 128, fused output", spec=dict(rows=128)), + Cell(why="1 row, fused output", spec=dict(rows=1)), + Cell(why="33 rows, contiguous gate|up (no extra columns)", spec=dict(rows=33, width=1280)), + Cell(why="zero rows", spec=dict(rows=0)), + ) + + note = """ + Metadata consumed: none; every row independent. Call site: the prefill target's MoE + block, on the shared expert's fused `[gate | up | gate logit | pad]` GEMM output. + + Preconditions: 2-D `x` with unit column stride, `x.shape[1] >= 2 * inter`; bf16 or fp16. + + Launch: `(rows, cdiv(inter, 1024))` programs of 4 warps. GB300, 16384 rows of the + 1288-wide output: 16.4 us (`flashinfer_silu_and_mul` on a contiguous copy of the gate|up + columns: 20.5 us for the kernel, 62.1 with the copy). + + Numerics: `g / (1 + exp(-g)) * u` in fp32 (Triton's `exp` and division), rounded once; + bitwise `flashinfer_silu_and_mul` on the same values (16384 x 640 random bf16 values, + the test's cells). + """ + + def __call__(self, x: torch.Tensor, inter: int) -> torch.Tensor: + rows = x.shape[0] + out = torch.empty((rows, inter), dtype=x.dtype, device=x.device) + if rows > 0: + _silu_and_mul_rows_kernel[(rows, triton.cdiv(inter, _BLOCK))]( + x, out, x.stride(0), INTER=inter, BLOCK=_BLOCK, num_warps=4 + ) + return out + + def reference(self, x: torch.Tensor, inter: int) -> torch.Tensor: + """fp32; the gate half first.""" + gate, up = x[:, : 2 * inter].float().chunk(2, dim=-1) + return (F.silu(gate) * up).to(x.dtype) + + def is_valid(self, x: torch.Tensor, inter: int) -> None: + assert x.dim() == 2 and x.stride(1) == 1, "x [T, width] with unit column stride" + assert x.shape[1] >= 2 * inter, "x carries the gate and up halves" + assert x.dtype in (torch.bfloat16, torch.float16), "bf16 / fp16" + + +silu_and_mul_rows = SiluAndMulRows() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/bf16_gemv_pdl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/bf16_gemv_pdl.py new file mode 100644 index 000000000000..2e9d3e382f15 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/bf16_gemv_pdl.py @@ -0,0 +1,392 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Decode-row bf16 projection launched with programmatic dependent launch (PDL), +its weight streaming in while the kernel before it finishes. + +Tier 5 (CuTe DSL written for this target). At decode row counts a projection +reads its weight once and its activation is a few KB, so the GEMV's time is the +weight stream plus its launch and ramp. A normal launch cannot read a byte of +that stream until the producer of the activation has exited. Launched with PDL, +every CTA loads its weight rows -- static, so safe to read early -- into +registers before `griddepcontrol.wait` and triggers its own dependents right +before it; the activation is read, and the output written, only after the wait. + +One CTA computes `OUTS` output columns over the whole K for every row; element +`k` of a row is held by thread `(k // VEC) % THREADS` as its vector +`(k // VEC) // THREADS`, so each k-tile of loads is contiguous across a warp. +Products and sums in fp32: a warp reduction, then a cross-warp one in shared +memory; one rounding at the store. + +A K with a shared-memory tier (`SOUTS`) gives each CTA `SOUTS` more outputs whose +weight rows it stages in shared memory with cp.async before the wait, each thread +exactly the vectors it later reads (so `cp.async.wait_group` alone orders them). The +2560 x 6144 projection then runs as one wave of 640 CTAs whose whole weight is on +chip by the wait, instead of 1280 CTAs of which the last third start loading theirs +only as the first ones exit. +""" + +from __future__ import annotations + +import functools +import os + +import cuda.bindings.driver as cuda_drv +import cutlass +import cutlass.cute as cute +import torch +from cutlass import const_expr +from cutlass._mlir.dialects import llvm +from cutlass.cute import experimental as cute_ext +from cutlass.cute.nvgpu import cpasync +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import dsl_user_op + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +#: Largest row count one launch serves; every row re-reads the weight from registers. +MAX_ROWS = 4 +#: (threads, outputs per CTA, elements per vector load) per K. Swept on GB300 at one row, +#: graph replay over per-layer weights: K=2560 64 threads x 2 outputs; K=640 (80 16-byte +#: vectors, which no warp multiple divides) 32 threads x 4 outputs on 8-byte vectors. +TILING = {2560: (64, 2, 8), 6144: (128, 2, 8), 640: (32, 4, 4)} +#: Outputs per CTA whose weight rows sit in shared memory (beside the `TILING` ones in +#: registers). K=6144 swept behind the GDN recurrence and the sparse-attention merge +#: (one row, graph replay): 2 + 2 a CTA beat 2 + 0, 2 + 6, 4 + 4 and 1 + 3. +SOUTS = {6144: 2} +_COMPILE_OPTIONS = "--ptxas-options -maxrregcount=64" +_COMPILE_OPTIONS_SMEM = "--ptxas-options -maxrregcount=80" +#: CTAs a multiprocessor must hold for a grid with a shared-memory tier: four, the +#: 640-CTA grid's share of 160 SMs. Behind the sparse-attention merge (one row, graph +#: replay) 4 runs the attention tail 13.5 us a layer against 14.6 at 1 and 15.6 at 2-3; +#: the GDN tail is unchanged. +MIN_BLOCKS_SMEM = 4 + + +@functools.lru_cache(maxsize=None) +def pdl_enabled(device_index: int) -> bool: + """Programmatic dependent launch: sm_90+ unless TRTLLM_ENABLE_PDL=0.""" + props = torch.cuda.get_device_properties(device_index) + return props.major >= 9 and os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" + + +def tiling(n: int, k: int) -> tuple[int, int, int, int]: + """`(threads, outs, souts, vec)` for a `[*, K] x [K, N]` product: register and shared-memory + outputs a CTA. Without a shared-memory tier that divides N the outputs shrink to divide N.""" + threads, outs, vec = TILING[k] + souts = SOUTS.get(k, 0) + if souts and n % (outs + souts) == 0: + return threads, outs, souts, vec + while n % outs: + outs //= 2 + return threads, outs, 0, vec + + +#: Before its PDL wait a CTA also has the weight rows of CTA `bidx + dist` brought into +#: L2: behind a latency-bound kernel only part of a large grid's CTAs fit, and the later +#: waves' weights would otherwise start streaming only after it. `dist` is about the CTAs +#: resident before the wait: 1216 behind the hyper-connection kernels (measured), at +#: least 912 (a 1280-CTA out projection's first wave) otherwise. Grids smaller than `dist` +#: prefetch nothing. +L2_DIST = 1216 +L2_DIST_MIN = 912 + + +def l2_dist(grid: int) -> int: + return min(L2_DIST, max(grid // 2, L2_DIST_MIN)) + + +@dsl_user_op +def _bulk_prefetch_l2(src, nbytes, *, loc=None, ip=None): + llvm.inline_asm( + None, + [ + cutlass.Int64(src.toint(loc=loc, ip=ip)).ir_value(loc=loc, ip=ip), + cutlass.Int32(nbytes).ir_value(loc=loc, ip=ip), + ], + "cp.async.bulk.prefetch.L2.global [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +class _GemvKernel: + def __init__(self, *, rows, k, threads, outs, vec, pdl, dist=L2_DIST, souts=0): + assert k % (threads * vec) == 0 and threads % 32 == 0 + self.rows, self.threads, self.outs, self.souts, self.vec = rows, threads, outs, souts, vec + self.dist = dist + self.k_tiles = k // (threads * vec) + self.warps = threads // 32 + self.pdl = pdl + + def smem_bytes(self) -> int: + red = self.rows * (self.outs + self.souts) * self.warps * 4 + if not self.souts: + return red + w = self.souts * self.k_tiles * self.threads * self.vec * 2 + return -(-w // 128) * 128 + -(-red // 128) * 128 + 128 + + @cute.jit + def __call__(self, mA, mW, mOut, stream: cuda_drv.CUstream): + n = cute.size(mW, mode=[0]) + load_a = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + cutlass.BFloat16, + num_bits_per_copy=16 * self.vec, + load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS, + ) + load_w = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + cutlass.BFloat16, + num_bits_per_copy=16 * self.vec, + load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING, + ) + g2s = cute.make_copy_atom( + cpasync.CopyG2SOp(), cutlass.BFloat16, num_bits_per_copy=16 * self.vec + ) + self.kernel(mA, mW, mOut, load_a, load_w, g2s).launch( + grid=[cute.ceil_div(n, self.outs + self.souts), 1, 1], + block=[self.threads, 1, 1], + smem=self.smem_bytes(), + stream=stream, + use_pdl=self.pdl, + min_blocks_per_mp=MIN_BLOCKS_SMEM if self.souts else 1, + ) + + @cute.kernel + def kernel( + self, mA, mW, mOut, load_a: cute.CopyAtom, load_w: cute.CopyAtom, g2s: cute.CopyAtom + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + warp = cute.arch.warp_idx() + rows: cutlass.Constexpr = self.rows + outs: cutlass.Constexpr = self.outs + threads: cutlass.Constexpr = self.threads + vec: cutlass.Constexpr = self.vec + k_tiles: cutlass.Constexpr = self.k_tiles + warps: cutlass.Constexpr = self.warps + souts: cutlass.Constexpr = self.souts + tot: cutlass.Constexpr = outs + souts + n0 = bidx * tot + + # (row, (vec, (threads, k_tiles))): thread t's vectors of a row, k-tile by k-tile. + w_vec = cute.logical_divide(cute.logical_divide(mW, (None, vec)), (None, (None, threads))) + a_vec = cute.logical_divide(cute.logical_divide(mA, (None, vec)), (None, (None, threads))) + smem = cutlass.utils.SmemAllocator() + + # The weight is static: it streams before the wait, while the producer runs. The + # shared-memory tier first ([output, k-tile, thread, vec]: a warp's vectors are + # consecutive), so it lands while the register tier streams. + if const_expr(souts > 0): + s_w = smem.allocate_tensor( + cutlass.BFloat16, + cute.make_layout( + (souts, k_tiles, threads, vec), + stride=(k_tiles * threads * vec, threads * vec, vec, 1), + ), + byte_alignment=128, + ) + for j in cutlass.range_constexpr(souts): + s_row = w_vec[n0 + outs + j, (None, (tidx, None))] + for kt in cutlass.range_constexpr(k_tiles): + cute.copy(g2s, s_row[None, kt], s_w[j, kt, tidx, None]) + cute.arch.cp_async_commit_group() + w_reg = cute.make_rmem_tensor( + cute.make_layout((outs, k_tiles, vec), stride=(k_tiles * vec, vec, 1)), cutlass.BFloat16 + ) + for j in cutlass.range_constexpr(outs): + w_row = w_vec[n0 + j, (None, (tidx, None))] + for kt in cutlass.range_constexpr(k_tiles): + cute.copy(load_w, w_row[None, kt], w_reg[j, kt, None]) + if const_expr(self.pdl): + n_rows = cute.size(mW, mode=[0]) + k_len = cute.size(mW, mode=[1]) + if tidx == 0: + if (bidx + self.dist + 1) * tot <= n_rows: + _bulk_prefetch_l2( + mW.iterator + (bidx + self.dist) * tot * k_len, tot * k_len * 2 + ) + cute.arch.griddepcontrol_launch_dependents() + cute.arch.griddepcontrol_wait() + + acc = cute.make_rmem_tensor(cute.make_layout((rows, tot), stride=(tot, 1)), cutlass.Float32) + acc.fill(0.0) + a_reg = cute.make_rmem_tensor( + cute.make_layout((k_tiles, vec), stride=(vec, 1)), cutlass.BFloat16 + ) + sw_reg = cute.make_rmem_tensor(cute.make_layout((vec,), stride=(1,)), cutlass.BFloat16) + for r in cutlass.range_constexpr(rows): + a_row = a_vec[r, (None, (tidx, None))] + for kt in cutlass.range_constexpr(k_tiles): + cute.copy(load_a, a_row[None, kt], a_reg[kt, None]) + for kt in cutlass.range_constexpr(k_tiles): + for v in cutlass.range_constexpr(vec): + av = a_reg[kt, v].to(cutlass.Float32) + for j in cutlass.range_constexpr(outs): + acc[r, j] = acc[r, j] + av * w_reg[j, kt, v].to(cutlass.Float32) + if const_expr(souts > 0): + if const_expr(r == 0): + cute.arch.cp_async_wait_group(0) + for j in cutlass.range_constexpr(souts): + for kt in cutlass.range_constexpr(k_tiles): + cute.autovec_copy(s_w[j, kt, tidx, None], sw_reg) + for v in cutlass.range_constexpr(vec): + acc[r, outs + j] = acc[r, outs + j] + a_reg[kt, v].to( + cutlass.Float32 + ) * sw_reg[v].to(cutlass.Float32) + for r in cutlass.range_constexpr(rows): + for j in cutlass.range_constexpr(tot): + acc[r, j] = cute.arch.warp_reduction_sum(acc[r, j]) + part = smem.allocate_tensor( + cutlass.Float32, + cute.make_layout((rows, tot, warps), stride=(tot * warps, warps, 1)), + byte_alignment=16, + ) + with cute.arch.elect_one(): + for r in cutlass.range_constexpr(rows): + for j in cutlass.range_constexpr(tot): + part[r, j, warp] = acc[r, j] + cute.arch.sync_threads() + if tidx < tot: + for r in cutlass.range_constexpr(rows): + total = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(warps): + total = total + part[r, tidx, w] + mOut[r, n0 + tidx] = total.to(cutlass.BFloat16) + + +@functools.lru_cache(maxsize=None) +def _compiled(rows, n, k, pdl): + threads, outs, souts, vec = tiling(n, k) + kern = _GemvKernel( + rows=rows, + k=k, + threads=threads, + outs=outs, + vec=vec, + pdl=pdl, + dist=l2_dist(-(-n // (outs + souts))), + souts=souts, + ) + a = torch.empty(rows, k, dtype=torch.bfloat16, device="cuda") + w = torch.empty(n, k, dtype=torch.bfloat16, device="cuda") + out = torch.empty(rows, n, dtype=torch.bfloat16, device="cuda") + stream = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream) + return cute_ext.compile( + kern, + from_dlpack(a, assumed_align=32), + from_dlpack(w, assumed_align=32), + from_dlpack(out, assumed_align=32), + stream, + options=_COMPILE_OPTIONS_SMEM if souts else _COMPILE_OPTIONS, + ) + + +def launch(mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool) -> torch.Tensor: + """The launch behind `Bf16GemvPdl.__call__`. `pdl` is honoured only where + `pdl_enabled` allows it.""" + m, k = mat_a.shape + n = mat_b.shape[1] + dev = mat_a.device + pdl = pdl and pdl_enabled(dev.index if dev.index is not None else torch.cuda.current_device()) + out = torch.empty((m, n), dtype=mat_a.dtype, device=dev) + if m > 0: + fn = _compiled(m, n, k, bool(pdl)) + stream = cuda_drv.CUstream(torch.cuda.current_stream(dev).cuda_stream) + fn( + from_dlpack(mat_a, assumed_align=32), + from_dlpack(mat_b.t(), assumed_align=32), + from_dlpack(out, assumed_align=32), + stream, + ) + return out + + +def _cell(why: str, m: int, n: int, k: int, pdl: bool = True) -> Cell: + return Cell(why=why, spec=dict(m=m, n=n, k=k, pdl=pdl)) + + +class Bf16GemvPdl(OpWrapper): + """`mat_a @ mat_b` in bf16 with fp32 accumulation, one PDL launch per call. + + A drop-in for `catalog/gemm/low_m_bf16_mm` at decode row counts: `mat_a` is + dense row-major `[M, K]` with `M <= MAX_ROWS`, `mat_b` the dense column-major + `[K, N]` view of a stored `[N, K]` weight (its `.t()`), the result a fresh + row-major `[M, N]` bf16 tensor. With `pdl` (and PDL available) the kernel is + launched as a programmatic dependent: it may start before the kernel ahead of + it on the stream exits, reads only its weight until `griddepcontrol.wait`, + and lets the kernel after it launch early as well. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # Every shape is a qwen3.8-flash-next-nvfp4 tp1 decode projection. + CELLS: tuple[Cell, ...] = ( + _cell("qwen3.8 decode M=1 router N=512", 1, 512, 2560), + _cell("qwen3.8 decode M=1 shared gate_up N=1280", 1, 1280, 2560), + _cell("qwen3.8 decode M=1 shared gate N=1", 1, 1, 2560), + _cell("qwen3.8 decode M=1 shared down K=640", 1, 2560, 640), + _cell("qwen3.8 decode M=3 shared down K=640, odd row count", 3, 2560, 640), + _cell("qwen3.8 decode M=1 GDN in-proj concat N=16480", 1, 16480, 2560), + _cell("qwen3.8 decode M=2 GDN in-proj concat", 2, 16480, 2560), + _cell("qwen3.8 decode M=1 GDN out / attention o, K=6144", 1, 2560, 6144), + _cell("qwen3.8 decode M=1 attention in-proj concat N=14080", 1, 14080, 2560), + _cell("qwen3.8 decode M=1 PLE key N=10240", 1, 10240, 2560), + _cell("qwen3.8 decode M=1 PLE value N=K=2560", 1, 2560, 2560), + _cell("qwen3.8 decode M=4 GDN out / attention o", 4, 2560, 6144), + _cell("qwen3.8 decode M=1 router, PDL off", 1, 512, 2560, pdl=False), + ) + + note = """ + Metadata consumed: none. Preconditions: bf16 operands; `mat_a` dense + row-major `[M, K]`, `1 <= M <= MAX_ROWS`; `mat_b` the dense column-major + `[K, N]` view of a `[N, K]` weight (stride `(1, K)`); K a `TILING` key. + + PDL: with `pdl` the launch carries the programmatic-dependent attribute. + Every read of `mat_a` and the store of the output come after + `griddepcontrol.wait`, so a producer launched before this kernel is + complete when they happen, whatever that producer is. Under a CUDA graph a + cross-stream dependency is an ordinary edge and stays a full dependency. + + Compilation: CuTe DSL, JIT on the first call per (M, N, K, pdl); that call + must happen before CUDA-graph capture -- the engine's warmup forward at each + captured batch size does it. + + Numerics: fp32 products and sums, one rounding to bf16; only the summation + order differs from the fp32 reference. + + Hard limits + M rows per launch boundary MAX_ROWS=4 driven at 1, 2, 3, 4 (cells) + K a TILING key 2560, 6144, 640 all driven + N any; the outputs per CTA halve until they divide N N=1 driven + """ + + def __call__(self, mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool = True) -> torch.Tensor: + return launch(mat_a, mat_b, pdl) + + def reference(self, mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool = True) -> torch.Tensor: + """fp32-accumulated, TF32 off; `pdl` is a launch attribute and is ignored.""" + with true_fp32_matmul(): + return (mat_a.float() @ mat_b.float()).to(mat_a.dtype) + + def is_valid(self, mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool = True) -> None: + assert mat_a.dtype == torch.bfloat16 and mat_b.dtype == torch.bfloat16, "bf16 operands only" + assert mat_a.dim() == 2 and mat_a.is_contiguous(), "mat_a must be dense row-major [M, K]" + assert 1 <= mat_a.shape[0] <= MAX_ROWS, f"M={mat_a.shape[0]} outside 1..{MAX_ROWS}" + assert mat_b.shape[0] == mat_a.shape[1], "inner dimensions differ" + assert mat_b.stride(0) == 1 and mat_b.stride(1) == mat_b.shape[0], ( + "mat_b must be dense column-major [K, N] (e.g. weight.t())" + ) + assert mat_a.shape[1] in TILING, f"K={mat_a.shape[1]} has no tiling" + + +bf16_gemv_pdl = Bf16GemvPdl() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_permuted.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_permuted.py new file mode 100644 index 000000000000..a925831d4631 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_permuted.py @@ -0,0 +1,144 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""NVFP4 MoE layer entered with router logits, left unfinalized: the trtllm-gen +runner's Renormalize routing + grouped FC1 GEMM + SwiGLU + NVFP4 +requantization + grouped FC2 GEMM, in one call, returning the permuted FC2 rows +with the routing weights and the expanded->permuted index map instead of the +combined `[T, H]` rows. The consumer forms the top-k weighted sum itself.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .fp4_block_scale_moe_renorm_routed import ACT_SWIGLU, ROUTING_RENORMALIZE, _ckpt +from .fp4_block_scale_moe_renorm_routed import fp4_block_scale_moe_renorm_routed as _finalized + + +def finalize( + rows: torch.Tensor, topk_weights: torch.Tensor, topk_rows: torch.Tensor +) -> torch.Tensor: + """The runner's finalize, emulated: per token `sum_k w[t, k] * rows[topk_rows[t, k]]` + as fp32 FMAs in slot order from zero (a -1 slot skipped), rounded once to the + rows' dtype. Each fp32 FMA is formed exactly in float64 (a product of two + fp32 values is exact there) and rounded to fp32. Returns `[T, H]`.""" + acc = torch.zeros(topk_rows.shape[0], rows.shape[1], dtype=torch.float32, device=rows.device) + for k in range(topk_rows.shape[1]): + idx = topk_rows[:, k].long() + r = rows[idx.clamp(min=0)].float().double() + w = topk_weights[:, k].float().double().unsqueeze(-1) + step = (r * w + acc.double()).float() + acc = torch.where((idx >= 0).unsqueeze(-1), step, acc) + return acc.to(rows.dtype) + + +class Fp4BlockScaleMoeRenormPermuted(OpWrapper): + """`fp4_block_scale_moe_renorm_routed` with `do_finalize=False`. + + Same operands, same routing, same GEMMs; the runner skips its finalize + launch and returns three tensors: + + - `rows` `[P, H]` bf16: the FC2 output in the runner's permuted, padded row + order (`P` is the runner's padded maximum; rows no slot maps to hold + garbage); + - `topk_weights` `[T, k]` bf16: the routing weights, slot order; + - `topk_rows` `[T, k]` int32: slot `(t, k)`'s row in `rows`, or -1. + + `finalize(rows, topk_weights, topk_rows)` (this module) is bitwise the + finalized entry's output on the same operands. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="_moe_block decode, finalize folded into the next site: one token", spec=_ckpt(1)), + Cell(why="4 rows: the last count on the block routing kernel", spec=_ckpt(4)), + Cell(why="5 rows: past the block kernel's row limit", spec=_ckpt(5)), + Cell(why="16 rows: decode_first's MAX_ROWS", spec=_ckpt(16)), + Cell(why="128 rows: the con=128 decode batch", spec=_ckpt(128)), + Cell(why="coarse 3-level logits: ties decide most slots", spec=_ckpt(64, dist="coarse")), + ) + + note = """ + Measured against the finalized entry on identical operands: the emulated + finalize of the three outputs is bitwise its output at every cell, so the + routing semantics, the GEMMs and the weight layout are the finalized + entry's (see its note). Which finalize kernel the runner would have + launched (scalar below ~1184 blocks, vectorized above) does not change the + arithmetic: both accumulate fp32 FMAs in slot order and round once. + + `rows` is the runner's own `[P, H]` buffer, allocated per call at the + padded maximum the runner computes from `T`, `k` and the expert count; + only rows some slot maps to hold data. `topk_weights` carries the output + dtype (bf16), as the finalize reads it. + + Silent, refused by `is_valid`: the finalized entry's refusals (a + non-contiguous operand, an uncertified (num_experts, top_k) pair, hidden + not a multiple of 256). + + Hard limits + T certified 1..128 (decode); the decode forward folds at + every site that takes the combine + num_experts 512, top_k 10 + """ + + def __call__( + self, + routing_logits: torch.Tensor, + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + output1_scale_scalar: torch.Tensor, + output1_scale_gate_scalar: torch.Tensor, + output2_scale_scalar: torch.Tensor, + num_experts: int, + top_k: int, + intermediate_size: int, + ) -> tuple: + rows, topk_weights, topk_rows = torch.ops.trtllm.fp4_block_scale_moe_runner( + routing_logits, + None, + hidden_states, + hidden_states_scale, + gemm1_weights, + gemm1_weights_scale, + None, + None, + None, + None, + gemm2_weights, + gemm2_weights_scale, + None, + output1_scale_scalar, + output1_scale_gate_scalar, + output2_scale_scalar, + num_experts, + top_k, + None, + None, + intermediate_size, + 0, + num_experts, + None, + ROUTING_RENORMALIZE, + False, + ACT_SWIGLU, + ) + return rows, topk_weights, topk_rows + + def reference(self, *args, **kwargs) -> torch.Tensor: + """The finalized entry's fp32 reference `[T, H]`: what + `finalize(*self(...))` must answer (same arguments as that entry's + `reference`).""" + return _finalized.reference(*args, **kwargs) + + def is_valid(self, *args, **kwargs) -> None: + _finalized.is_valid(*args, **kwargs) + + +fp4_block_scale_moe_renorm_permuted = Fp4BlockScaleMoeRenormPermuted() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_routed.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_routed.py new file mode 100644 index 000000000000..f74817c05d6a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_routed.py @@ -0,0 +1,225 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""NVFP4 MoE layer entered with router logits: the trtllm-gen runner's own +Renormalize routing (top-k over the raw logits, softmax over those k) + grouped +FC1 GEMM + SwiGLU + NVFP4 requantization + grouped FC2 GEMM + routing-weighted +combine, in one trtllm-gen call.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fp4_block_scale_moe_runner import ( # noqa: E501 + fp4_block_scale_moe_runner as _pre_routed, +) + +from .renorm_moe_routing_op import renorm_moe_routing_op as _renorm + +#: `RoutingMethodType::Renormalize`: NoOp preprocess, top-k, softmax over the k. +ROUTING_RENORMALIZE = 1 +#: `act_type` SwiGlu. +ACT_SWIGLU = 0 + +#: The (num_experts, top_k) pairs this entry certifies. The Renormalize routing +#: policy dispatches on a compiled tier list, and a pair outside every tier is a +#: logged error with no exception and **no routing launch** -- the GEMMs then +#: read stale routing buffers. So the guard admits what was measured, not what +#: the tier list might also cover. +CERTIFIED_PAIRS = frozenset({(512, 10)}) + + +def _ckpt(tokens: int, dist: str = "randn", logits_dtype=torch.bfloat16) -> dict: + """qwen4_exp's routed geometry: 512 experts, top-10, H 2560 / I 640.""" + return dict( + tokens=tokens, + hidden=2560, + intermediate=640, + num_experts=512, + top_k=10, + logits_dtype=logits_dtype, + dist=dist, + ) + + +class Fp4BlockScaleMoeRenormRouted(OpWrapper): + """One NVFP4 MoE layer whose routing happens inside the call. + + `routing_logits` is the router GEMM's `[T, num_experts]` output. Per token + the kernel picks the `top_k` largest logits (ties to the LOWER expert index) + and weights them by an fp32 softmax over just those `top_k`, rounded to + bf16 -- the same function as `ops/moe/renorm_moe_routing_op` followed by the + pre-routed `catalog/moe/fp4_block_scale_moe_runner`, and bitwise equal to + that composition at every cell. Returns the combined `[T, H]` bf16 rows. + + Fixed, not arguments: no routing bias, Renormalize routing, SwiGlu, no + expert parallelism (`local_num_experts == num_experts`), `do_finalize=True`. + The weight layout (row shuffle + 128x4 scale swizzle) and the three + per-expert scalars are exactly the pre-routed entry's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="_moe_block decode: one token", spec=_ckpt(1)), + Cell(why="4 rows: the last count on the block routing kernel", spec=_ckpt(4)), + Cell(why="5 rows: past the block kernel's row limit", spec=_ckpt(5)), + Cell(why="8 rows: a decode CUDA-graph batch", spec=_ckpt(8)), + Cell(why="32 rows: a decode CUDA-graph batch", spec=_ckpt(32)), + Cell(why="128 rows: the con=128 decode batch", spec=_ckpt(128)), + Cell(why="1024 rows: a prefill chunk", spec=_ckpt(1024)), + Cell(why="8192 rows: max_num_tokens of the identity config", spec=_ckpt(8192)), + Cell(why="16384 rows: max_num_tokens of the tuned config", spec=_ckpt(16384)), + Cell(why="coarse 3-level logits: ties decide most slots", spec=_ckpt(64, dist="coarse")), + Cell( + why="fp32 logits: the runner's other accepted logits dtype", + spec=_ckpt(64, logits_dtype=torch.float32), + ), + ) + + note = """ + Routing semantics, measured against the composition it replaces + (`renorm_moe_routing_op` -> pre-routed runner): bitwise-equal finalized + output at every cell, randn / coarse / large-magnitude / all-equal logits, + bf16 and fp32 logits, 1..16384 rows. Both kernels rank by a packed + (twiddled value, 65535 - index) key (ties to the lower expert index), take + an fp32 `expf` softmax over the k with the same 32-lane warp reduction, and + round the weight to bf16 once. + + Kernel choice inside the call is by row count (block kernel up to a small + count, single-cluster up to the cluster cap, then the multi-kernel + histogram path); every one of them shares the top-k and softmax code. + `routing_bias` is None (the thop defaults the bias dtype to bf16 and never + reads it on this policy). PDL (`TRTLLM_ENABLE_PDL`) is read by the routing + launch. + + Silent, refused by `is_valid`: a non-contiguous operand (raw pointers, + dense layout assumed) and a (num_experts, top_k) pair outside + `CERTIFIED_PAIRS` (a tier-list miss logs and launches nothing). + + Hard limits + T certified 1..16384 (max_num_tokens of the tuned config). + Autotuner buckets stop at 8192, the op's + `tune_max_num_tokens` default, which this entry does not + pass: a call past 8192 rows runs the 8192 bucket's tactic + num_experts 512, top_k 10 (CERTIFIED_PAIRS) + """ + + def __call__( + self, + routing_logits: torch.Tensor, + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + output1_scale_scalar: torch.Tensor, + output1_scale_gate_scalar: torch.Tensor, + output2_scale_scalar: torch.Tensor, + num_experts: int, + top_k: int, + intermediate_size: int, + ) -> torch.Tensor: + return torch.ops.trtllm.fp4_block_scale_moe_runner( + routing_logits, + None, + hidden_states, + hidden_states_scale, + gemm1_weights, + gemm1_weights_scale, + None, + None, + None, + None, + gemm2_weights, + gemm2_weights_scale, + None, + output1_scale_scalar, + output1_scale_gate_scalar, + output2_scale_scalar, + num_experts, + top_k, + None, + None, + intermediate_size, + 0, + num_experts, + None, + ROUTING_RENORMALIZE, + True, + ACT_SWIGLU, + )[0] + + def reference( + self, + routing_logits: torch.Tensor, + hidden_states: torch.Tensor, + top_k: int, + up: tuple, + gate: tuple, + down: tuple, + fc2_global_scale: float, + ) -> torch.Tensor: + """The routing, then the layer, both by their own entries' references. + + Routing: `renorm_moe_routing_op.reference` (stable descending sort in the + logits' dtype, fp32 softmax over the k, bf16 weights -- the dtype the + runner's routing writes). The layer: the pre-routed runner entry's fp32 + reference over the dequantized `(codes, scales)` per-expert operands, as + that entry documents (`hidden_states` is the exact fp32 value of the + NVFP4 activation). Returns fp32 `[T, H]`. + """ + ids, weights = _renorm.reference(routing_logits, top_k, torch.bfloat16) + return _pre_routed.reference(hidden_states, ids, weights, up, gate, down, fc2_global_scale) + + def is_valid( + self, + routing_logits: torch.Tensor, + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + output1_scale_scalar: torch.Tensor, + output1_scale_gate_scalar: torch.Tensor, + output2_scale_scalar: torch.Tensor, + num_experts: int, + top_k: int, + intermediate_size: int, + ) -> None: + # Raw data pointers, dense row-major assumed: a strided view is read as + # if dense. A router GEMM output written into a wider buffer is the + # realistic way to get here. + for name, tensor in ( + ("routing_logits", routing_logits), + ("hidden_states", hidden_states), + ("hidden_states_scale", hidden_states_scale), + ("gemm1_weights", gemm1_weights), + ("gemm1_weights_scale", gemm1_weights_scale), + ("gemm2_weights", gemm2_weights), + ("gemm2_weights_scale", gemm2_weights_scale), + ("output1_scale_scalar", output1_scale_scalar), + ("output1_scale_gate_scalar", output1_scale_gate_scalar), + ("output2_scale_scalar", output2_scale_scalar), + ): + assert tensor.is_contiguous(), ( + f"{name} must be contiguous; a strided view is read as if dense " + "and silently produces wrong results" + ) + # A tier-list miss in the Renormalize routing policy logs an error and + # launches nothing; the GEMMs then consume stale routing buffers. + assert (num_experts, top_k) in CERTIFIED_PAIRS, ( + f"(num_experts, top_k) = ({num_experts}, {top_k}) is not a certified " + f"Renormalize routing pair {sorted(CERTIFIED_PAIRS)}" + ) + # Same layout rule as the pre-routed entry: the prepared block-scale + # layout is silently misread below a hidden multiple of 256. + assert gemm1_weights.shape[-1] * 2 % 256 == 0, ( + "hidden size (gemm1_weights.shape[-1] * 2) must be a multiple of 256" + ) + + +fp4_block_scale_moe_renorm_routed = Fp4BlockScaleMoeRenormRouted() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_moe_renorm_cute_dsl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_moe_renorm_cute_dsl.py new file mode 100644 index 000000000000..a02b4e47257b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_moe_renorm_cute_dsl.py @@ -0,0 +1,220 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""NVFP4 MoE layer entered with router logits, on TRT-LLM's CuTe DSL Blackwell +MoE kernels: Renormalize routing (`renorm_moe_routing_op`), the token sort, +a gathered grouped FC1 GEMM with SwiGLU and NVFP4 requantization in its +epilogue, and a grouped FC2 GEMM that scatter-adds the routing-weighted rows +into the output -- the layer `fp4_block_scale_moe_renorm_routed` computes, for +prefill-sized row counts, from a second copy of the experts in this kernel +family's layout.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .fp4_block_scale_moe_renorm_routed import CERTIFIED_PAIRS, _ckpt +from .fp4_block_scale_moe_renorm_routed import fp4_block_scale_moe_renorm_routed as _routed +from .renorm_moe_routing_op import renorm_moe_routing_op + +#: Token tile of the grouped GEMMs (the kernels' M tile). +TILE_TOKENS = 128 +_FP4 = torch.float4_e2m1fn_x2 + + +class Fp4MoeRenormCuteDsl(OpWrapper): + """One NVFP4 MoE layer whose routing happens inside the call, on the CuTe + DSL MoE kernels. Same function as `fp4_block_scale_moe_renorm_routed`: + per token the `top_k` largest logits (ties to the LOWER expert index), + weighted by an fp32 softmax over just those `top_k` rounded to bf16; the + experts' SwiGLU FC1, the NVFP4 requantization of its output, FC2, and the + weighted sum. Returns the combined `[T, H]` bf16 rows. + + Operands, in this kernel family's layout (see `weights.py`): `fc1_w` + `[E, 2I, H/2]` uint8, `[up ; gate]` rows interleaved in 64-row groups; + `fc1_s` `[E, 2I, H/16]` e4m3 bytes in that row order, 128x4-swizzled; + `fc2_w` `[E, H, I/2]` as stored; `fc2_s` `[E, H, I/16]` swizzled. + `hidden_states` / `hidden_states_scale` are the NVFP4 activation with + **linear** block scales, quantized with the layer's FC1 activation global + scale. Scalars: `fc1_alpha` `[E]` dequantizes the FC1 accumulator (the + routed entry's `output1_scale_gate_scalar`); `fc2_input_scale` `[1]` is the + FC2 activation's quantization multiplier (the reciprocal of its dequant + scale), one number for every expert; `fc2_alpha` `[E]` dequantizes the FC2 + accumulator (the routed entry's `output2_scale_scalar`). + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="1024 rows: a prefill chunk", spec=_ckpt(1024)), + Cell(why="4096 rows: where the target switches to this entry", spec=_ckpt(4096)), + Cell(why="8192 rows", spec=_ckpt(8192)), + Cell( + why="16377 rows: a con=128 mixed step (16256 context + 121 generation)", + spec=_ckpt(16377), + ), + Cell(why="16384 rows: max_num_tokens of the tuned config", spec=_ckpt(16384)), + ) + + note = """ + Semantics against `fp4_block_scale_moe_renorm_routed` (trtllm-gen) on the + same experts: the same routing (both are the `renorm_moe_routing_op` + function, weights rounded to bf16 once) and the same NVFP4 arithmetic per + expert (e2m1 x e4m3-block x fp32-alpha products accumulated in fp32, SwiGLU + on the FC1 accumulator, the FC2 input requantized to NVFP4 with 16-element + blocks), but not bitwise: the GEMMs' accumulation order differs, and the + weighted sum over a token's experts is a scatter-add with fp32 atomics, so + its order -- and the output's last bit -- varies between launches. Both sit + inside the routed entry's band against its fp32 reference. + + One FC2 activation scale for all experts: the routed entry takes it per + expert (folded into `output1_scale_scalar`); a checkpoint whose experts + disagree on `down_proj.input_scale` cannot use this entry as is. + + Kernels (all TRT-LLM's): `trtllm::moe_sort` (tile metadata), + `trtllm::moe_output_memset_inplace` (zeroes the rows the scatter-add + touches), `trtllm::cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell`, + `trtllm::cute_dsl_nvfp4_grouped_gemm_finalize_inplace_blackwell`; the two + GEMMs pick their tactics through the AutoTuner. Not CUDA-graph certified + (prefill steps run eagerly). + + Hard limits + T certified 1024..16384 + num_experts 512, top_k 10 (the routed entry's CERTIFIED_PAIRS) + """ + + def __call__( + self, + routing_logits: torch.Tensor, + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + fc1_w: torch.Tensor, + fc1_s: torch.Tensor, + fc1_alpha: torch.Tensor, + fc2_w: torch.Tensor, + fc2_s: torch.Tensor, + fc2_alpha: torch.Tensor, + fc2_input_scale: torch.Tensor, + num_experts: int, + top_k: int, + ) -> torch.Tensor: + tokens = hidden_states.shape[0] + hidden = hidden_states.shape[1] * 2 + ids, weights = renorm_moe_routing_op(routing_logits, top_k, torch.bfloat16) + scales = weights.float() + t2e, t2lim, e2p, p2e, _, live = torch.ops.trtllm.moe_sort( + token_selected_experts=ids, + token_final_scales=scales, + num_experts=num_experts, + top_k=top_k, + local_expert_offset=0, + local_num_experts=num_experts, + tile_tokens_dim=TILE_TOKENS, + ) + out = torch.empty((tokens, hidden), dtype=torch.bfloat16, device=hidden_states.device) + torch.ops.trtllm.moe_output_memset_inplace( + input=out, + tile_idx_to_mn_limit=t2lim, + expanded_idx_to_permuted_idx=e2p, + permuted_idx_to_expanded_idx=p2e, + num_non_exiting_tiles=live, + tile_tokens_dim=TILE_TOKENS, + top_k=top_k, + ep_size=1, + enable_alltoall=False, + ) + h, h_sf = torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell( + input=hidden_states.view(_FP4), + weight=fc1_w.view(_FP4), + input_scale=hidden_states_scale.view(torch.uint8).view(tokens, hidden // 16), + weight_scale=fc1_s.view(torch.uint8), + alpha=fc1_alpha, + tile_idx_to_group_idx=t2e, + tile_idx_to_mn_limit=t2lim, + permuted_idx_to_expanded_idx=p2e, + num_non_exiting_tiles=live, + global_sf=fc2_input_scale, + num_experts=num_experts, + top_k=top_k, + num_local_experts=num_experts, + local_expert_offset=0, + tile_size=TILE_TOKENS, + ) + torch.ops.trtllm.cute_dsl_nvfp4_grouped_gemm_finalize_inplace_blackwell( + input=h.view(_FP4), + weight=fc2_w.view(_FP4), + input_scale=h_sf.view(torch.uint8), + weight_scale=fc2_s.view(torch.uint8), + alpha=fc2_alpha, + output=out, + tile_idx_to_group_idx=t2e, + tile_idx_to_mn_limit=t2lim, + permuted_idx_to_expanded_idx=p2e, + num_non_exiting_tiles=live, + token_final_scales=scales, + num_experts=num_experts, + top_k=top_k, + num_local_experts=num_experts, + local_expert_offset=0, + tile_size=TILE_TOKENS, + output_dtype=torch.bfloat16, + ) + return out + + def reference( + self, + routing_logits: torch.Tensor, + hidden_states: torch.Tensor, + top_k: int, + up: tuple, + gate: tuple, + down: tuple, + fc2_global_scale: float, + ) -> torch.Tensor: + """The routed entry's fp32 reference (routing, then the layer over the + dequantized `(codes, scales)` per-expert operands); fp32 `[T, H]`.""" + return _routed.reference( + routing_logits, hidden_states, top_k, up, gate, down, fc2_global_scale + ) + + def is_valid( + self, + routing_logits: torch.Tensor, + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + fc1_w: torch.Tensor, + fc1_s: torch.Tensor, + fc1_alpha: torch.Tensor, + fc2_w: torch.Tensor, + fc2_s: torch.Tensor, + fc2_alpha: torch.Tensor, + fc2_input_scale: torch.Tensor, + num_experts: int, + top_k: int, + ) -> None: + for name, tensor in ( + ("routing_logits", routing_logits), + ("hidden_states", hidden_states), + ("hidden_states_scale", hidden_states_scale), + ("fc1_w", fc1_w), + ("fc1_s", fc1_s), + ("fc2_w", fc2_w), + ("fc2_s", fc2_s), + ): + assert tensor.is_contiguous(), f"{name} must be contiguous" + assert (num_experts, top_k) in CERTIFIED_PAIRS, ( + f"(num_experts, top_k) = ({num_experts}, {top_k}) is not certified" + ) + tokens, half_hidden = hidden_states.shape + assert hidden_states_scale.numel() == tokens * half_hidden * 2 // 16, ( + "hidden_states_scale must hold one linear e4m3 scale a 16-element block" + ) + assert fc1_alpha.shape == (num_experts,) and fc2_alpha.shape == (num_experts,), ( + "per-expert alphas" + ) + assert fc2_input_scale.numel() == 1, "one FC2 activation scale for every expert" + + +fp4_moe_renorm_cute_dsl = Fp4MoeRenormCuteDsl() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp8_block_scale_moe_runner.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp8_block_scale_moe_runner.py index fcb1f0eac64a..9fe9ee74e6a4 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp8_block_scale_moe_runner.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp8_block_scale_moe_runner.py @@ -122,9 +122,13 @@ class Fp8BlockScaleMoeRunner(OpWrapper): 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", + why="8192 tokens: max_num_tokens of the identity config", spec=dict(tokens=8192, **_CKPT), ), + Cell( + why="16384 rows: max_num_tokens of the tuned config", + spec=dict(tokens=16384, **_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), @@ -217,7 +221,11 @@ class Fp8BlockScaleMoeRunner(OpWrapper): states): an upstream defect, eager is unaffected. Hard limits - T certified 1..8192 (max_num_tokens); chunk above it, not probed + T certified 1..16384 (max_num_tokens of the tuned config); + chunk above it, not probed. Autotuner buckets stop at + `tune_max_num_tokens` (default 8192, which the target + leaves): past it a call runs the 8192 bucket's tactic; a + larger value makes the op append a bucket at that size top_k <= 32 (raises past it); certified 1, 2, 6, 8, 10, 16 E, H, I certified up to 512, 7168, 2048 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/renorm_moe_routing_op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/renorm_moe_routing_op.py index 3e100830050c..261f6660ecbf 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/renorm_moe_routing_op.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/renorm_moe_routing_op.py @@ -56,6 +56,7 @@ class RenormMoeRoutingOp(OpWrapper): spec=_ckpt(4096), ), Cell(why="8192 tokens: max_num_tokens", spec=_ckpt(8192)), + Cell(why="16384 rows: max_num_tokens of the tuned config", spec=_ckpt(16384)), Cell( why="the default fp32 output dtype at the checkpoint router", spec=_ckpt(512, torch.float32), @@ -124,7 +125,7 @@ class RenormMoeRoutingOp(OpWrapper): Hard limits num_experts <= 512 (raises past it) topk <= 16 (raises past it); > num_experts silent, guarded - T certified 0..8192 (16384 measured once) + T certified 0..16384 (max_num_tokens of the tuned config) """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_gate_conv_inject_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_gate_conv_inject_norm.py new file mode 100644 index 000000000000..9a7ef343b849 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_gate_conv_inject_norm.py @@ -0,0 +1,676 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""PLE block tail for decode-shaped batches: the key and query grouped norms, +the per-stream sigmoid gate, the conv-input grouped norm, the dilated +depthwise short conv against the paged conv state (published back), the +residual add into the hyper-connection stream, and the next site's grouped +norm -- three Triton launches behind one call. + +Tier 5 (Triton written for this target). It replaces, on the decode path, +three `_grouped_norm` calls (12 ``flashinfer_gemma_rmsnorm`` + 3 cats), +``gated_ngram_injection`` (~10 launches), ``dilated_short_conv_paged`` (~15 +launches around ``F.conv1d``), two adds and the norm-only +``qwen4_exp_hc_inject_norm`` after the PLE add. No ``torch.ops.trtllm`` op or +FlashInfer callable takes a dilation (see ``dilated_short_conv_paged``). + +The three launches split where the parallelism changes shape. The norms and +the gate are row reductions, one program per (row, stream). The conv and the +state publish are per channel, and each sequence's conv state is 184 KB of +`[channels, 9]` rows: one program per (sequence, block of channels), so the +state moves as contiguous spans across many SMs instead of strided through +four. The residual add and the next site's norm are row reductions again. +Folding all three into one (row, stream) program was measured 2.5-3x slower +at 128 rows: four programs per sequence cannot move the state fast enough. +""" + +from __future__ import annotations + +import math +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .ple_ngram_embed import PER_SEQ_MAX_ROWS, _seq_of_rows + + +@triton.jit +def _rms(x, w, HIDDEN: tl.constexpr, eps): + """Gemma RMS norm in fp32: `x * rsqrt(mean(x^2) + eps) * (1 + w)`.""" + return x * tl.rsqrt(tl.sum(x * x, axis=0) / HIDDEN + eps) * (1.0 + w) + + +@triton.jit +def _ple_gate_kernel( + hc_ptr, + key_ptr, + value_ptr, + nk_ptr, + nq_ptr, + nc_ptr, + gated_ptr, + xc_ptr, + stride_hc, + stride_key, + stride_value, + stride_out, + eps, + clamp_min, + sqrt_d, + HIDDEN: tl.constexpr, + BLOCK: tl.constexpr, +): + """Per (row, stream): the key and query grouped norms, the gate, the gated + value and the conv-input grouped norm. Writes `gated` and `xc`.""" + row = tl.program_id(0).to(tl.int64) + stream = tl.program_id(1) + cols = tl.arange(0, BLOCK) + m = cols < HIDDEN + ch = stream * HIDDEN + cols + x = tl.load(hc_ptr + row * stride_hc + ch, mask=m, other=0.0).to(tl.float32) + kp = tl.load(key_ptr + row * stride_key + ch, mask=m, other=0.0).to(tl.float32) + v = tl.load(value_ptr + row * stride_value + cols, mask=m, other=0.0).to(tl.float32) + wk = tl.load(nk_ptr + ch, mask=m, other=0.0).to(tl.float32) + wq = tl.load(nq_ptr + ch, mask=m, other=0.0).to(tl.float32) + wc = tl.load(nc_ptr + ch, mask=m, other=0.0).to(tl.float32) + kn = _rms(kp, wk, HIDDEN, eps).to(tl.bfloat16).to(tl.float32) + qn = _rms(x, wq, HIDDEN, eps).to(tl.bfloat16).to(tl.float32) + g = tl.sum(kn * qn, axis=0) / sqrt_d + sg = tl.sqrt(tl.maximum(tl.abs(g), clamp_min)) + sg = tl.where(g > 0, sg, tl.where(g < 0, -sg, 0.0)) + gated = ((1.0 / (1.0 + tl.exp(-sg))) * v).to(tl.bfloat16) + xc = _rms(gated.to(tl.float32), wc, HIDDEN, eps).to(tl.bfloat16) + tl.store(gated_ptr + row * stride_out + ch, gated, mask=m) + tl.store(xc_ptr + row * stride_out + ch, xc, mask=m) + + +@triton.jit +def _ple_conv_kernel( + xc_ptr, + gated_ptr, + ple_ptr, + cw_ptr, + state_ptr, + slots_ptr, + carry_ptr, + qsl_ptr, + stride_out, + stride_state, + pad_slot, + KSIZE: tl.constexpr, + DIL: tl.constexpr, + STATE: tl.constexpr, + HEAD: tl.constexpr, + MAX_P: tl.constexpr, + ROWS: tl.constexpr, + HAS_CARRY: tl.constexpr, +): + """Per (sequence, block of ROWS channels): the dilated depthwise conv of + each of the sequence's rows against `[state | conv inputs]`, its SiLU added + to the gated value (`ple`), and the state advanced in place. + + A channel's state is `STATE` contiguous columns. The publish moves it as a + `[ROWS, HEAD]` tile of the first HEAD (a power of two) columns plus, when + `STATE == HEAD + 1`, one `[ROWS]` vector for the last, so no lane is spent + on padding and every index is a plain 2-D offset. Every load is issued + before the first store.""" + s = tl.program_id(0) + ch0 = tl.program_id(1) * ROWS + begin = tl.load(qsl_ptr + s).to(tl.int64) + length = (tl.load(qsl_ptr + s + 1).to(tl.int64) - begin).to(tl.int32) + slot = tl.load(slots_ptr + s).to(tl.int64) + live = slot != pad_slot + # A sequence without carry reads a zero state (it still publishes). + if HAS_CARRY: + seeded = live & (tl.load(carry_ptr + s) != 0) + else: + seeded = live + slot = tl.where(live, slot, 0) + block = state_ptr + slot * stride_state + ch0 * STATE + row0 = begin * stride_out + ch0 + xrows = xc_ptr + row0 + r = tl.arange(0, ROWS) + i = tl.arange(0, HEAD) + + # The published state, `window[length : length + STATE]` per channel: the + # surviving state columns, then this sequence's newest conv inputs. + src = length + i + keep = tl.load( + block + r[:, None] * STATE + src[None, :], mask=seeded & (src < STATE)[None, :], other=0.0 + ) + fresh = tl.load( + xrows + (src - STATE)[None, :] * stride_out + r[:, None], + mask=live & (src >= STATE)[None, :], + other=0.0, + ) + head = tl.where((src < STATE)[None, :], keep, fresh) + if STATE > HEAD: + src_t = length + HEAD + keep_t = tl.load(block + r * STATE + src_t, mask=seeded & (src_t < STATE), other=0.0) + fresh_t = tl.load( + xrows + (src_t - STATE) * stride_out + r, mask=live & (src_t >= STATE), other=0.0 + ) + tail = tl.where(src_t < STATE, keep_t, fresh_t) + + ch = ch0 + r + for j in tl.static_range(MAX_P): + valid = j < length + gated = tl.load(gated_ptr + row0 + j * stride_out + r, mask=valid, other=0.0) + acc = tl.zeros([ROWS], dtype=tl.float32) + # Tap k of row j reads window column j + k * DIL of `[state | xc rows]`. + for k in tl.static_range(KSIZE): + q = j + k * DIL + wtap = tl.load(cw_ptr + ch * KSIZE + k).to(tl.float32) + if q < STATE: + tap = tl.load(block + r * STATE + q, mask=seeded, other=0.0) + else: + tap = tl.load(xrows + (q - STATE) * stride_out + r, mask=valid, other=0.0) + acc += wtap * tap.to(tl.float32) + conv = acc.to(tl.bfloat16).to(tl.float32) + y = (conv / (1.0 + tl.exp(-conv))).to(tl.bfloat16).to(tl.float32) + out = (gated.to(tl.float32) + y).to(tl.bfloat16) + tl.store(ple_ptr + row0 + j * stride_out + r, out, mask=valid) + + # Every read of these state rows -- `keep` and the taps -- lands before + # the first store over them. + tl.debug_barrier() + tl.store(block + r[:, None] * STATE + i[None, :], head, mask=live) + if STATE > HEAD: + tl.store(block + r * STATE + HEAD, tail, mask=live) + + +@triton.jit +def _ple_conv_rows_kernel( + xc_ptr, + gated_ptr, + ple_ptr, + cw_ptr, + state_ptr, + slots_ptr, + carry_ptr, + qsl_ptr, + num_rows, + num_seqs, + stride_out, + stride_state, + pad_slot, + KSIZE: tl.constexpr, + DIL: tl.constexpr, + STATE: tl.constexpr, + BLOCK_T: tl.constexpr, + ROWS: tl.constexpr, + LOG_S: tl.constexpr, + HAS_CARRY: tl.constexpr, +): + """Context-shaped batches: per (BLOCK_T packed rows, ROWS channels), each + row's dilated conv over `[state | conv inputs]` of its own sequence, its + SiLU added to the gated value. Reads the state only; the publish is the + per-sequence kernel with no rows, launched after.""" + t = tl.program_id(0).to(tl.int64) * BLOCK_T + tl.arange(0, BLOCK_T) + tmask = t < num_rows + ch = tl.program_id(1) * ROWS + tl.arange(0, ROWS) + s = _seq_of_rows(qsl_ptr, t, num_seqs, LOG_S) + begin = tl.load(qsl_ptr + s, mask=tmask, other=0).to(tl.int64) + tpos = t - begin + slot = tl.load(slots_ptr + s, mask=tmask, other=pad_slot).to(tl.int64) + seeded = slot != pad_slot + if HAS_CARRY: + seeded = seeded & (tl.load(carry_ptr + s, mask=tmask, other=0) != 0) + slot = tl.where(seeded, slot, 0) + m2 = tmask[:, None] + offs = t[:, None] * stride_out + ch[None, :] + gated = tl.load(gated_ptr + offs, mask=m2, other=0.0) + acc = tl.zeros([BLOCK_T, ROWS], dtype=tl.float32) + for k in tl.static_range(KSIZE): + wtap = tl.load(cw_ptr + ch * KSIZE + k).to(tl.float32) + q = tpos + k * DIL + from_state = q < STATE + tap_s = tl.load( + state_ptr + slot[:, None] * stride_state + ch[None, :] * STATE + q[:, None], + mask=m2 & (from_state & seeded)[:, None], + other=0.0, + ) + tap_x = tl.load( + xc_ptr + offs + (k * DIL - STATE) * stride_out, + mask=m2 & (~from_state)[:, None], + other=0.0, + ) + tap = tl.where(from_state[:, None], tap_s, tap_x) + acc += wtap[None, :] * tap.to(tl.float32) + conv = acc.to(tl.bfloat16).to(tl.float32) + y = (conv / (1.0 + tl.exp(-conv))).to(tl.bfloat16).to(tl.float32) + out = (gated.to(tl.float32) + y).to(tl.bfloat16) + tl.store(ple_ptr + offs, out, mask=m2) + + +@triton.jit +def _ple_residual_norm_kernel( + hc_ptr, + ple_ptr, + nx_ptr, + hc_out_ptr, + normed_ptr, + stride_hc, + stride_ple, + stride_hc_out, + stride_normed, + eps, + HIDDEN: tl.constexpr, + BLOCK: tl.constexpr, +): + """Per (row, stream): the PLE output added into the stream, then the next + site's grouped norm of the result.""" + row = tl.program_id(0).to(tl.int64) + stream = tl.program_id(1) + cols = tl.arange(0, BLOCK) + m = cols < HIDDEN + ch = stream * HIDDEN + cols + x = tl.load(hc_ptr + row * stride_hc + ch, mask=m, other=0.0).to(tl.float32) + ple = tl.load(ple_ptr + row * stride_ple + ch, mask=m, other=0.0).to(tl.float32) + wx = tl.load(nx_ptr + ch, mask=m, other=0.0).to(tl.float32) + hn = (x + ple).to(tl.bfloat16) + tl.store(hc_out_ptr + row * stride_hc_out + ch, hn, mask=m) + nrm = _rms(hn.to(tl.float32), wx, HIDDEN, eps) + tl.store(normed_ptr + row * stride_normed + ch, nrm.to(tl.bfloat16), mask=m) + + +def _state_head(state_len: int) -> int: + """The largest power of two not above `state_len`: the conv kernel's tile + width; the one column past it, if any, moves as a vector.""" + return 1 << (state_len.bit_length() - 1) + + +#: Channels per conv program. +_CONV_ROWS = 128 +_CONV_WARPS = 4 +#: Context-shaped conv tiles: packed rows x channels per program. +_ROWS_BLOCK_T = 32 +_ROWS_CHANNELS = 128 +_ROWS_WARPS = 4 +_MIN_LOG_S = 8 + +#: This checkpoint's PLE tail geometry: 4 streams of 2560, conv kernel 4 at +#: dilation 3 (a 9-column state), the gate floor `_PLE_CLAMP_MIN`. +CKPT_HC_COUNT = 4 +CKPT_HIDDEN = 2560 +CKPT_DILATION = 3 +CKPT_CLAMP_MIN = 1e-6 +#: Largest PLE-conv pool seen for this checkpoint (see `dilated_short_conv_paged`). +ENGINE_SLOTS = 461 + + +def _cell(why: str, lengths, **spec) -> Cell: + return Cell(why=why, spec=dict(lengths=tuple(lengths), **spec)) + + +class PleGateConvInjectNorm(OpWrapper): + """With `C = hc_count`, `D = value.shape[-1]`, every grouped norm a gemma + RMS norm over each `D`-wide stream (`n(x, w) = x * rsqrt(mean(x^2) + eps) + * (1 + w)`), and each step rounded to the activation dtype where the + composition it replaces rounds -- for row `t` of sequence `s`, stream `c`:: + + kn = round(n(key[t, c], key_norm[c])) + qn = round(n(hc[t, c], query_norm[c])) + g = sum(kn * qn) / sqrt(D); sg = sign(g) * sqrt(max(|g|, clamp_min)) + gt = round(sigmoid(sg) * value[t]) # gated_ngram_injection + xc = round(n(gt, conv_norm[c])) + conv = round(sum_k conv_weight[c, k] * window[c, t' + k * dilation]) + ple = round(gt + round(silu(conv))) # dilated_short_conv_paged + hc_out[t, c] = round(hc[t, c] + ple) + normed[t, c] = round(n(hc_out[t, c], next_norm[c])) # hc_inject_norm, norm only + + where `window` is `[conv_states[slot] | xc of the sequence's rows]` and + `t'` the row's offset in its sequence; the state part is zero for a + sequence whose `has_initial_states` flag is false (None: all carry). Then `conv_states[slot]` becomes + `window[:, L : L + state_len]`, the raw conv-input tail. Every sequence + carries its state (a decode batch); a `pad_slot_id` sequence reads a zero + state and publishes nothing. + + `max_rows > PER_SEQ_MAX_ROWS` (context-shaped steps) takes a fourth + launch: the conv over packed (row tile, channel block) programs, the + per-sequence kernel then only publishing the state. + + Returns `(hc_out, normed)`, fresh `[T, C*D]` in `hc`'s dtype. + `conv_states` is advanced in place. + + Fusion boundary: everything from the key/query norms to the next site's + norm. The n-gram embedding and the key / value projections are the + caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: lengths (rows per sequence), pool_slots (default num_seqs + 2), + # slots (None -> scattered distinct), max_rows (max(lengths)), pad (index + # of a sequence to give pad_slot_id), hc_count (4), hidden (2560), dtype + # (bf16), layout "contiguous" | "row_slice" (hc and key are buf[::2]). + CELLS: tuple[Cell, ...] = tuple( + _cell(f"decode {n} x 1 row at the 461-slot pool", (1,) * n, pool_slots=ENGINE_SLOTS) + for n in (1, 8, 32, 128) + ) + ( + _cell("decode 256 x 1 row, 512-slot pool", (1,) * 256, pool_slots=512), + _cell("MTP verify p=2: 8 x 2 rows", (2,) * 8, pool_slots=ENGINE_SLOTS), + _cell("MTP verify p=4: 32 x 4 rows", (4,) * 32, pool_slots=ENGINE_SLOTS), + _cell("MTP verify p=3: 4 x 3 rows", (3,) * 4), + _cell("a pad_slot_id row beside live rows", (1,) * 6, pad=2), + _cell("rows shorter than max_rows, a zero-row sequence", (2, 0, 1), max_rows=4), + _cell("hc and key row slices buf[::2]", (1,) * 16, layout="row_slice"), + _cell("hc_count 2 at hidden 2560", (1,) * 8, hc_count=2), + _cell("hidden 1024 (not the checkpoint)", (1,) * 8, hidden=1024), + # Context-shaped (prefill / mixed) steps: the packed-row conv. + _cell("prefill 8192 rows, fresh", (8192,), carries=(False,), pool_slots=ENGINE_SLOTS), + _cell( + "chunked prefill 16384 rows carrying", + (16384,), + carries=(True,), + pool_slots=ENGINE_SLOTS, + ), + _cell( + "mixed step: 2 contexts + 64 decode rows", + (3001, 5000) + (1,) * 64, + carries=(False, True) + (True,) * 64, + pool_slots=ENGINE_SLOTS, + ), + _cell( + "short contexts, tile-straddling, zero-row, pad", + (40, 1, 0, 7, 33, 2), + carries=(True, False, True, False, True, True), + pad=4, + ), + _cell("mixed step, carry None", (9, 1, 1)), + ) + + note = """ + Metadata consumed: none. Call site: the decode target's PLE block, once + per forward: `hc` the `[T, 10240]` stream, `key` the key projection + `[T, 10240]` and `value` the value projection `[T, 2560]` (both cuBLAS, + bf16), the three PLE norm weights and the next site's attention-side norm + weight as contiguous `[10240]`, the `[10240, 1, 4]` conv weight, the whole + `[slots, 10240, 9]` conv pool, int32 slots, int64 `query_start_loc_long` + (the rebuild under MTP verify), `max_rows = gen_p`. + + Preconditions (data; not guarded): at most `max_rows` rows per sequence; + distinct live slots; `query_start_loc[-1] == T`. Operands 2-D with unit + column stride (row strides free); weights contiguous. + + Numerics: every rounding point of the composition is kept, so the only + differences are fp32 reduction order (the three norms' sums of squares + and the gate's dot) and fp32 `exp` -- each can move a value across a + rounding midpoint by one ulp, and that ulp propagates down the chain. + + CUDA-graph capturable: no host sync, outputs and scratch (`gated` / `ple`, + `xc`) allocated per call. Three launches: gate (rows x streams), conv + (sequences x channel blocks of `_CONV_ROWS`), residual + norm (rows x + streams). + """ + + def __call__( + self, + hc: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + key_norm: torch.Tensor, + query_norm: torch.Tensor, + conv_norm: torch.Tensor, + conv_weight: torch.Tensor, + conv_states: torch.Tensor, + state_indices: torch.Tensor, + query_start_loc: torch.Tensor, + next_norm: torch.Tensor, + eps: float, + hc_count: int, + clamp_min: float, + dilation: int, + pad_slot_id: int, + max_rows: int, + has_initial_states: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + rows, width = hc.shape + hidden = width // hc_count + ksize = conv_weight.shape[-1] + state_len = conv_states.shape[-1] + hc_out = torch.empty((rows, width), dtype=hc.dtype, device=hc.device) + normed = torch.empty((rows, width), dtype=hc.dtype, device=hc.device) + num_seqs = state_indices.shape[0] + if num_seqs == 0 or rows == 0: + return hc_out, normed + gated = torch.empty((rows, width), dtype=hc.dtype, device=hc.device) + xc = torch.empty((rows, width), dtype=hc.dtype, device=hc.device) + block = triton.next_power_of_2(hidden) + warps = 8 if block >= 2048 else 4 + _ple_gate_kernel[(rows, hc_count)]( + hc, + key, + value, + key_norm, + query_norm, + conv_norm, + gated, + xc, + hc.stride(0), + key.stride(0), + value.stride(0), + gated.stride(0), + eps, + clamp_min, + math.sqrt(hidden), + HIDDEN=hidden, + BLOCK=block, + num_warps=warps, + ) + # `ple` reuses `gated`'s storage row for row: each element is read by + # the program that then writes it. + conv_rows = min(_CONV_ROWS, width) + has_carry = has_initial_states is not None + carry = has_initial_states if has_carry else state_indices + per_seq = max_rows <= PER_SEQ_MAX_ROWS + if not per_seq: + tile_ch = min(_ROWS_CHANNELS, width) + _ple_conv_rows_kernel[(triton.cdiv(rows, _ROWS_BLOCK_T), width // tile_ch)]( + xc, + gated, + gated, + conv_weight, + conv_states, + state_indices, + carry, + query_start_loc, + rows, + num_seqs, + gated.stride(0), + conv_states.stride(0), + pad_slot_id, + KSIZE=ksize, + DIL=dilation, + STATE=state_len, + BLOCK_T=_ROWS_BLOCK_T, + ROWS=tile_ch, + LOG_S=max(_MIN_LOG_S, (num_seqs - 1).bit_length()), + HAS_CARRY=has_carry, + num_warps=_ROWS_WARPS, + ) + # Per sequence: the rows (decode) and the state publish. On the + # context-shaped path it carries no rows and only publishes, after + # every program above has read the state. + _ple_conv_kernel[(num_seqs, width // conv_rows)]( + xc, + gated, + gated, + conv_weight, + conv_states, + state_indices, + carry, + query_start_loc, + gated.stride(0), + conv_states.stride(0), + pad_slot_id, + KSIZE=ksize, + DIL=dilation, + STATE=state_len, + HEAD=_state_head(state_len), + MAX_P=max_rows if per_seq else 0, + ROWS=conv_rows, + HAS_CARRY=has_carry, + num_warps=_CONV_WARPS, + ) + _ple_residual_norm_kernel[(rows, hc_count)]( + hc, + gated, + next_norm, + hc_out, + normed, + hc.stride(0), + gated.stride(0), + hc_out.stride(0), + normed.stride(0), + eps, + HIDDEN=hidden, + BLOCK=block, + num_warps=warps, + ) + return hc_out, normed + + def reference( + self, + hc: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + key_norm: torch.Tensor, + query_norm: torch.Tensor, + conv_norm: torch.Tensor, + conv_weight: torch.Tensor, + conv_states: torch.Tensor, + state_indices: torch.Tensor, + query_start_loc: torch.Tensor, + next_norm: torch.Tensor, + eps: float, + hc_count: int, + clamp_min: float, + dilation: int, + pad_slot_id: int, + max_rows: int, + has_initial_states: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """`(hc_out, normed, conv_states)` as the call leaves them: float64 + throughout, rounded to the activation dtype exactly where the op (and + the composition it replaces) rounds; the conv one sequence at a time as + an explicit tap sum. Mutates nothing.""" + dtype = hc.dtype + rows, width = hc.shape + hidden = width // hc_count + + def rnd(x): + return x.to(dtype).double() + + def norm(x, w): + w = w.double().view(hc_count, hidden) + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * (1.0 + w) + + streams = hc.double().view(rows, hc_count, hidden) + kn = rnd(norm(key.double().view(rows, hc_count, hidden), key_norm)) + qn = rnd(norm(streams, query_norm)) + g = (kn * qn).sum(-1, keepdim=True) / math.sqrt(hidden) + sg = g.abs().clamp_min(clamp_min).sqrt() * g.sign() + gated = rnd(torch.sigmoid(sg) * value.double().unsqueeze(-2)) + xc = rnd(norm(gated, conv_norm)).reshape(rows, width) + + ksize = conv_weight.shape[-1] + state_len = conv_states.shape[-1] + w = conv_weight.double().reshape(width, ksize) + conv = torch.zeros(rows, width, dtype=torch.float64, device=hc.device) + pool = conv_states.clone() + starts = query_start_loc.tolist() + carries = ( + [True] * state_indices.shape[0] + if has_initial_states is None + else [bool(c) for c in has_initial_states.tolist()] + ) + for s, slot in enumerate(state_indices.tolist()): + begin, end = starts[s], starts[s + 1] + live = slot != pad_slot_id + prefix = ( + conv_states[slot].double() + if live and carries[s] + else torch.zeros(width, state_len, dtype=torch.float64, device=hc.device) + ) + window = torch.cat([prefix, xc[begin:end].t()], dim=1) + n = end - begin + acc = torch.zeros(width, n, dtype=torch.float64, device=hc.device) + for k in range(ksize): + acc += w[:, k : k + 1] * window[:, k * dilation : k * dilation + n] + conv[begin:end] = rnd(acc.t()) + if live: + pool[slot] = window[:, end - begin : end - begin + state_len].to(dtype) + ple = rnd(gated.reshape(rows, width) + rnd(torch.nn.functional.silu(conv))) + hc_out = (hc.double() + ple).to(dtype) + normed = norm(hc_out.double().view(rows, hc_count, hidden), next_norm) + return hc_out, normed.reshape(rows, width).to(dtype), pool + + def is_valid( + self, + hc: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + key_norm: torch.Tensor, + query_norm: torch.Tensor, + conv_norm: torch.Tensor, + conv_weight: torch.Tensor, + conv_states: torch.Tensor, + state_indices: torch.Tensor, + query_start_loc: torch.Tensor, + next_norm: torch.Tensor, + eps: float, + hc_count: int, + clamp_min: float, + dilation: int, + pad_slot_id: int, + max_rows: int, + has_initial_states: Optional[torch.Tensor] = None, + ) -> None: + rows, width = hc.shape + assert width % hc_count == 0, f"width {width} not divisible by hc_count {hc_count}" + hidden = width // hc_count + for name, t in (("hc", hc), ("key", key), ("value", value)): + assert t.dim() == 2 and t.stride(1) == 1, f"{name} must be 2-D, unit column stride" + assert key.shape == hc.shape, f"key {tuple(key.shape)} must match hc {tuple(hc.shape)}" + assert value.shape == (rows, hidden), f"value must be [{rows}, {hidden}]" + for name, t in ( + ("key_norm", key_norm), + ("query_norm", query_norm), + ("conv_norm", conv_norm), + ("next_norm", next_norm), + ): + assert t.shape == (width,) and t.is_contiguous(), f"{name} must be contiguous [{width}]" + assert conv_weight.dim() == 3 and conv_weight.shape[:2] == (width, 1), ( + "conv_weight [W, 1, K]" + ) + assert conv_weight.is_contiguous(), "conv_weight contiguous" + ksize = conv_weight.shape[-1] + assert conv_states.dim() == 3 and conv_states.shape[1:] == ( + width, + (ksize - 1) * dilation, + ), f"conv_states must be [slots, {width}, {(ksize - 1) * dilation}]" + assert conv_states.stride(2) == 1 and conv_states.stride(1) == conv_states.shape[2], ( + "conv_states rows must be dense" + ) + assert conv_states.dtype == hc.dtype, "conv state dtype must match the activation" + assert query_start_loc.shape[0] == state_indices.shape[0] + 1, ( + "query_start_loc must be [num_seqs + 1]" + ) + state_len = conv_states.shape[-1] + assert state_len - _state_head(state_len) <= 1, ( + f"a {state_len}-column conv state is not a power of two plus at most one" + ) + assert width % min(_CONV_ROWS, width) == 0, ( + f"width {width} must be a multiple of the conv block ({_CONV_ROWS} channels)" + ) + assert max_rows >= 1 + assert has_initial_states is None or has_initial_states.shape == state_indices.shape, ( + "has_initial_states must be one flag per sequence" + ) + + +ple_gate_conv_inject_norm = PleGateConvInjectNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_ngram_embed.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_ngram_embed.py new file mode 100644 index 000000000000..27d3a5e80837 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_ngram_embed.py @@ -0,0 +1,591 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""PLE n-gram embedding for decode-shaped batches: id history, n-gram hash, +table gather and per-tensor fp8 dequantization in one Triton launch, with the +per-slot n-gram context published back in the same pass. + +Tier 5 (Triton written for this target). It replaces, on the decode path, +the composition ``ngram_token_history`` (~20 launches) + ``ngram_hash`` (~30 +launches) + ``nn.functional.embedding`` + ``dequantize_e4m3_per_tensor`` -- +nothing registered under ``torch.ops.trtllm`` or in the installed FlashInfer hashes +n-grams or gathers-and-dequantizes, and each of those entries is launch-bound +at decode (0.3-0.4 ms eager each). + +One program per **sequence**, which loops over that sequence's rows: the carry +read and the publish of one slot happen in one program, so no other program can +observe a half-published row. + +Context-shaped batches (more than ``PER_SEQ_MAX_ROWS`` rows per sequence) take +a second form: one program per 32 packed rows, each row binary-searching +``query_start_loc`` for its sequence, then a one-program-per-sequence publish +launched after it -- two launches, bitwise equal to the first form. +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ..embedding.ngram_hash import ngram_hash +from .ngram_token_history import ngram_token_history + +#: This checkpoint's PLE geometry: 3-gram, 16 heads of 160 fp8 columns. +CKPT_NGRAM_SIZE = 3 +CKPT_HEAD_DIM = 160 +CKPT_TABLE_ROWS = 320001536 +CKPT_EOS_TOKEN_ID = 248044 +CKPT_SCALE = 0.00019931793212890625 +#: Largest `ple_ngram` pool seen for this checkpoint (see `ngram_token_history`). +ENGINE_SLOTS = 461 +#: Above this many rows per sequence (MTP verify carries at most 4) the call +#: takes the context-shaped path: packed row tiles plus a separate publish. +PER_SEQ_MAX_ROWS = 4 +_ROWS_BLOCK_T = 32 +#: The row-to-sequence search covers at least 2^8 sequences, so one compiled +#: variant serves every batch up to that size. +_MIN_LOG_S = 8 + + +@triton.jit +def _ple_ngram_embed_kernel( + ids_ptr, + pool_ptr, + slots_ptr, + carry_ptr, + qsl_ptr, + table_ptr, + scale_ptr, + mult_ptr, + vocab_ptr, + offs_ptr, + out_ptr, + stride_pool, + stride_table, + stride_out, + eos, + NGRAM: tl.constexpr, + HEADS_PER: tl.constexpr, + HEADS: tl.constexpr, + BLOCK_H: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_P: tl.constexpr, + BLOCK_C: tl.constexpr, + HAS_CARRY: tl.constexpr, +): + CTX: tl.constexpr = NGRAM - 1 + s = tl.program_id(0) + begin = tl.load(qsl_ptr + s).to(tl.int64) + length = tl.load(qsl_ptr + s + 1).to(tl.int64) - begin + slot = tl.load(slots_ptr + s).to(tl.int64) + pool_row = pool_ptr + slot * stride_pool + eos64 = tl.full([BLOCK_P], eos, tl.int64) + if HAS_CARRY: + fresh = tl.load(carry_ptr + s) == 0 + + t = tl.arange(0, BLOCK_P) + tmask = t < length + heads = tl.arange(0, BLOCK_H) + hmask = heads < HEADS + order = heads // HEADS_PER + 2 + vocab = tl.load(vocab_ptr + heads, mask=hmask, other=1) + offs = tl.load(offs_ptr + heads, mask=hmask, other=0) + + # Shift 0 is the token itself; shift k reads `k` back in the sequence's own + # stream, i.e. the packed ids or, before the slice starts, the carry row. + # Once a shift lands on an EOS, every deeper shift reads EOS: the hash's + # segment starts one past the most recent EOS strictly before the token. + mix = tl.load(ids_ptr + begin + t, mask=tmask, other=0).to(tl.int64) * tl.load(mult_ptr) + mixed = tl.zeros([BLOCK_P, BLOCK_H], dtype=tl.int64) + hit = t < 0 + for k in tl.static_range(1, NGRAM): + pos = t - k + from_ids = tl.load(ids_ptr + begin + pos, mask=tmask & (pos >= 0), other=0).to(tl.int64) + from_pool = tl.load(pool_row + (CTX + pos), mask=tmask & (pos < 0), other=0) + if HAS_CARRY: + from_pool = tl.where(fresh, eos64, from_pool) + ck = tl.where(pos >= 0, from_ids, from_pool) + sk = tl.where(hit, eos64, ck) + hit = hit | (ck == eos64) + mix = mix ^ (sk * tl.load(mult_ptr + k)) + mixed = tl.where((order == k + 1)[None, :], mix[:, None], mixed) + + rem = mixed % vocab[None, :] + rows = tl.where(rem < 0, rem + vocab[None, :], rem) + offs[None, :] + + d = tl.arange(0, BLOCK_D) + m3 = tmask[:, None, None] & hmask[None, :, None] & (d < HEAD_DIM)[None, None, :] + q = tl.load(table_ptr + rows[:, :, None] * stride_table + d[None, None, :], mask=m3, other=0.0) + val = q.to(tl.float32) * tl.load(scale_ptr).to(tl.float32) + cols = heads[None, :, None] * HEAD_DIM + d[None, None, :] + tl.store( + out_ptr + (begin + t)[:, None, None] * stride_out + cols, + val.to(out_ptr.dtype.element_ty), + mask=m3, + ) + + # Publish the last CTX ids of `carry + slice` -- after every thread's carry + # read, since the row being overwritten is the one just read. + c = tl.arange(0, BLOCK_C) + cmask = c < CTX + src = length + c + from_slice = tl.load(ids_ptr + begin + src - CTX, mask=cmask & (src >= CTX), other=0).to( + tl.int64 + ) + old = tl.load(pool_row + src, mask=cmask & (src < CTX), other=0) + if HAS_CARRY: + old = tl.where(fresh, tl.full([BLOCK_C], eos, tl.int64), old) + published = tl.where(src >= CTX, from_slice, old) + tl.debug_barrier() + tl.store(pool_row + c, published, mask=cmask) + + +@triton.jit +def _seq_of_rows(qsl_ptr, t, num_seqs, LOG_S: tl.constexpr): + """Per packed row `t`: the largest `s < num_seqs` with `qsl[s] <= t`, i.e. + the sequence that owns the row (zero-row sequences are skipped over).""" + s = tl.zeros_like(t) + for i in tl.static_range(LOG_S): + cand = s + (1 << (LOG_S - 1 - i)) + ok = cand < num_seqs + start = tl.load(qsl_ptr + cand, mask=ok, other=0) + s = tl.where(ok & (start <= t), cand, s) + return s + + +@triton.jit +def _ple_ngram_embed_rows_kernel( + ids_ptr, + pool_ptr, + slots_ptr, + carry_ptr, + qsl_ptr, + table_ptr, + scale_ptr, + mult_ptr, + vocab_ptr, + offs_ptr, + out_ptr, + num_rows, + num_seqs, + stride_pool, + stride_table, + stride_out, + eos, + NGRAM: tl.constexpr, + HEADS_PER: tl.constexpr, + HEADS: tl.constexpr, + BLOCK_H: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_T: tl.constexpr, + LOG_S: tl.constexpr, + HAS_CARRY: tl.constexpr, +): + """Context-shaped batches: one program per BLOCK_T packed rows, each row + finding its sequence in `query_start_loc`. Window, hash, gather and + dequantize only; the publish is `_ple_ngram_publish_kernel`, launched + after, because the rows that read a slot's carry may span two programs.""" + CTX: tl.constexpr = NGRAM - 1 + t = tl.program_id(0).to(tl.int64) * BLOCK_T + tl.arange(0, BLOCK_T) + tmask = t < num_rows + s = _seq_of_rows(qsl_ptr, t, num_seqs, LOG_S) + begin = tl.load(qsl_ptr + s, mask=tmask, other=0).to(tl.int64) + tpos = t - begin + slot = tl.load(slots_ptr + s, mask=tmask, other=0).to(tl.int64) + eos64 = tl.full([BLOCK_T], eos, tl.int64) + if HAS_CARRY: + fresh = tl.load(carry_ptr + s, mask=tmask, other=1) == 0 + + heads = tl.arange(0, BLOCK_H) + hmask = heads < HEADS + order = heads // HEADS_PER + 2 + vocab = tl.load(vocab_ptr + heads, mask=hmask, other=1) + offs = tl.load(offs_ptr + heads, mask=hmask, other=0) + + mix = tl.load(ids_ptr + t, mask=tmask, other=0).to(tl.int64) * tl.load(mult_ptr) + mixed = tl.zeros([BLOCK_T, BLOCK_H], dtype=tl.int64) + hit = tmask & (t < 0) + for k in tl.static_range(1, NGRAM): + pos = tpos - k + from_ids = tl.load(ids_ptr + t - k, mask=tmask & (pos >= 0), other=0).to(tl.int64) + from_pool = tl.load( + pool_ptr + slot * stride_pool + (CTX + pos), mask=tmask & (pos < 0), other=0 + ) + if HAS_CARRY: + from_pool = tl.where(fresh, eos64, from_pool) + ck = tl.where(pos >= 0, from_ids, from_pool) + sk = tl.where(hit, eos64, ck) + hit = hit | (ck == eos64) + mix = mix ^ (sk * tl.load(mult_ptr + k)) + mixed = tl.where((order == k + 1)[None, :], mix[:, None], mixed) + + rem = mixed % vocab[None, :] + rows = tl.where(rem < 0, rem + vocab[None, :], rem) + offs[None, :] + + d = tl.arange(0, BLOCK_D) + dmask = d < HEAD_DIM + m2 = tmask[:, None] & dmask[None, :] + scale = tl.load(scale_ptr).to(tl.float32) + for h in tl.static_range(HEADS): + row_h = tl.sum(tl.where((heads == h)[None, :], rows, 0), axis=1) + q = tl.load(table_ptr + row_h[:, None] * stride_table + d[None, :], mask=m2, other=0.0) + val = q.to(tl.float32) * scale + tl.store( + out_ptr + t[:, None] * stride_out + (h * HEAD_DIM + d)[None, :], + val.to(out_ptr.dtype.element_ty), + mask=m2, + ) + + +@triton.jit +def _ple_ngram_publish_kernel( + ids_ptr, + pool_ptr, + slots_ptr, + carry_ptr, + qsl_ptr, + stride_pool, + eos, + NGRAM: tl.constexpr, + BLOCK_C: tl.constexpr, + HAS_CARRY: tl.constexpr, +): + """One program per sequence: the last `NGRAM - 1` ids of `carry + slice` + published to its slot, after every row's carry read (stream order).""" + CTX: tl.constexpr = NGRAM - 1 + s = tl.program_id(0) + begin = tl.load(qsl_ptr + s).to(tl.int64) + length = tl.load(qsl_ptr + s + 1).to(tl.int64) - begin + slot = tl.load(slots_ptr + s).to(tl.int64) + pool_row = pool_ptr + slot * stride_pool + c = tl.arange(0, BLOCK_C) + cmask = c < CTX + src = length + c + from_slice = tl.load(ids_ptr + begin + src - CTX, mask=cmask & (src >= CTX), other=0).to( + tl.int64 + ) + old = tl.load(pool_row + src, mask=cmask & (src < CTX), other=0) + if HAS_CARRY: + fresh = tl.load(carry_ptr + s) == 0 + old = tl.where(fresh, tl.full([BLOCK_C], eos, tl.int64), old) + published = tl.where(src >= CTX, from_slice, old) + tl.debug_barrier() + tl.store(pool_row + c, published, mask=cmask) + + +def _cell(why: str, lengths, **spec) -> Cell: + return Cell(why=why, spec=dict(lengths=tuple(lengths), **spec)) + + +class PleNgramEmbed(OpWrapper): + """For sequence `s` over packed rows `[b, e)` with `slot = state_indices[s]`, + its id stream is `carry + input_ids[b:e]`, where `carry` is + `context_pool[slot]` (oldest first) when `has_initial_states[s]` -- or + always, when `has_initial_states` is None -- and `ngram_size - 1` copies of + `eos_token_id` when not. Row `t` of that slice gets the window + `stream[t : t + ngram_size]` (newest last), exactly `ngram_token_history`'s; + each window hashes to one global row per head exactly as `ngram_hash` with + `num_query_tokens=1`; and:: + + out[t, h * D : (h + 1) * D] = to(scale.dtype)(float(table[row[t, h]]) * float(scale)) + + with `D = table.shape[1]` -- `dequantize_e4m3_per_tensor`'s one rounding. + Then, **in place**, `context_pool[slot]` becomes the last `ngram_size - 1` + ids of the stream (a zero-row sequence republishes its carry). Unnamed slots + are untouched. + + Returns a fresh contiguous `[num_tokens, heads * D]` in `scale`'s dtype. + `context_pool` is the only argument modified. + + Fusion boundary: carry read and seed, window, hash, gather, dequantize, + publish. The key/value projections that consume the result are the caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: lengths (rows per sequence), slots (None -> arange), pool_slots + # (default num_seqs + 2), carries (None -> the argument is None), max_rows + # (default max(lengths)), eos_density (0.0), table "small" | "host" | + # "ckpt" (the full 320001536-row device table, rows filled where touched), + # ngram_size (3), heads_per_ngram (8), ids_dtype (int32). + CELLS: tuple[Cell, ...] = ( + tuple( + _cell(f"decode {n} x 1 row at the 461-slot pool", (1,) * n, pool_slots=ENGINE_SLOTS) + for n in (1, 8, 32, 128) + ) + + tuple( + _cell( + f"decode {n} x 1 row, host-resident table, 461 slots", + (1,) * n, + pool_slots=ENGINE_SLOTS, + table="host", + ) + for n in (1, 128) + ) + + ( + _cell( + "decode 128 x 1 row, the full checkpoint-height table", + (1,) * 128, + pool_slots=ENGINE_SLOTS, + table="ckpt", + eos_density=0.05, + ), + _cell("decode 64 x 1 row, EOS in stream and carry", (1,) * 64, eos_density=0.3), + _cell("MTP verify p=2: 8 x 2 rows", (2,) * 8, pool_slots=ENGINE_SLOTS), + _cell("MTP verify p=3: 8 x 3 rows", (3,) * 8, eos_density=0.2), + _cell("MTP verify p=4: 32 x 4 rows", (4,) * 32, pool_slots=ENGINE_SLOTS), + _cell( + "carry flags mixed, fresh rows seed EOS", + (1, 3, 1, 2), + carries=(True, False, True, False), + ), + _cell( + "zero-row sequences republish carry", + (2, 0, 1, 0), + carries=(True, True, True, False), + max_rows=2, + ), + _cell("rows shorter than max_rows", (1, 3, 2), max_rows=4), + _cell("int64 ids", (1,) * 16, ids_dtype=torch.int64), + _cell("ngram_size 4, 4 heads per order", (1, 2, 4), ngram_size=4, heads_per_ngram=4), + _cell("ngram_size 2, 8 heads", (1,) * 8, ngram_size=2, eos_density=0.2), + # Context-shaped (prefill / mixed) steps: the packed-row form. + _cell("prefill 8192 rows, fresh", (8192,), carries=(False,), pool_slots=ENGINE_SLOTS), + _cell( + "chunked prefill 16384 rows carrying, host table", + (16384,), + carries=(True,), + pool_slots=ENGINE_SLOTS, + table="host", + eos_density=0.01, + ), + _cell( + "mixed step: 2 contexts + 64 decode rows", + (3001, 5000) + (1,) * 64, + carries=(False, True) + (True,) * 64, + pool_slots=ENGINE_SLOTS, + eos_density=0.05, + ), + _cell( + "short contexts, tile-straddling, zero-row", + (40, 1, 0, 7, 33, 2), + carries=(True, False, True, False, True, True), + eos_density=0.2, + ), + _cell("mixed step, carry None", (9, 1, 1), max_rows=9), + ) + ) + + note = """ + Metadata consumed: none -- batch layout, pool, slot map are explicit. + Call site: the decode target's PLE block, once per forward: + `input_ids` int32 `[T]`, the whole `[slots, 2]` int64 n-gram pool, + int32 `state_indices[:batch]`, `has_initial_states=None` (every decode + sequence carries), int64 `query_start_loc_long` (or the + `cu_seqlens_from_seq_lens` rebuild under MTP verify), the + `[320001536, 160]` fp8 table (host-resident by default), the `[1, 1, 1]` + bf16 scale, `max_rows = gen_p`. + + Preconditions (data; not guarded -- each would need a sync): + - every sequence has at most `max_rows` rows; rows past it are neither + written nor published; + - slots real (`0 <= slot < slots`) and distinct; -1 is not supported; + - `query_start_loc[-1] == num_tokens`; + - token id * max multiplier < 2^63 (int64 wraps, as in `ngram_hash`). + + Bitwise equal to `ngram_token_history` + `ngram_hash` + embedding + + `dequantize_e4m3_per_tensor`, window, rows, values and pool alike. + CUDA-graph capturable: no host sync, one launch, output allocated per call. + Reads the table through its device address, so a host-mapped table needs + no copy. + """ + + def __call__( + self, + input_ids: torch.Tensor, + context_pool: torch.Tensor, + state_indices: torch.Tensor, + has_initial_states: Optional[torch.Tensor], + query_start_loc: torch.Tensor, + table: torch.Tensor, + scale: torch.Tensor, + layer_multipliers: torch.Tensor, + head_vocab_sizes: torch.Tensor, + head_offsets: torch.Tensor, + eos_token_id: int, + max_rows: int, + ) -> torch.Tensor: + num_tokens = input_ids.shape[0] + num_seqs = state_indices.shape[0] + ngram_size = layer_multipliers.numel() + heads = head_vocab_sizes.numel() + head_dim = table.shape[1] + out = torch.empty( + (num_tokens, heads * head_dim), dtype=scale.dtype, device=input_ids.device + ) + if num_seqs > 0 and max_rows > PER_SEQ_MAX_ROWS: + carry = has_initial_states if has_initial_states is not None else state_indices + if num_tokens > 0: + _ple_ngram_embed_rows_kernel[(triton.cdiv(num_tokens, _ROWS_BLOCK_T),)]( + input_ids, + context_pool, + state_indices, + carry, + query_start_loc, + table, + scale, + layer_multipliers, + head_vocab_sizes, + head_offsets, + out, + num_tokens, + num_seqs, + context_pool.stride(0), + table.stride(0), + out.stride(0), + eos_token_id, + NGRAM=ngram_size, + HEADS_PER=heads // (ngram_size - 1), + HEADS=heads, + BLOCK_H=triton.next_power_of_2(heads), + HEAD_DIM=head_dim, + BLOCK_D=triton.next_power_of_2(head_dim), + BLOCK_T=_ROWS_BLOCK_T, + LOG_S=max(_MIN_LOG_S, (num_seqs - 1).bit_length()), + HAS_CARRY=has_initial_states is not None, + num_warps=4, + ) + _ple_ngram_publish_kernel[(num_seqs,)]( + input_ids, + context_pool, + state_indices, + carry, + query_start_loc, + context_pool.stride(0), + eos_token_id, + NGRAM=ngram_size, + BLOCK_C=triton.next_power_of_2(ngram_size - 1), + HAS_CARRY=has_initial_states is not None, + num_warps=1, + ) + elif num_seqs > 0: + _ple_ngram_embed_kernel[(num_seqs,)]( + input_ids, + context_pool, + state_indices, + has_initial_states if has_initial_states is not None else state_indices, + query_start_loc, + table, + scale, + layer_multipliers, + head_vocab_sizes, + head_offsets, + out, + context_pool.stride(0), + table.stride(0), + out.stride(0), + eos_token_id, + NGRAM=ngram_size, + HEADS_PER=heads // (ngram_size - 1), + HEADS=heads, + BLOCK_H=triton.next_power_of_2(heads), + HEAD_DIM=head_dim, + BLOCK_D=triton.next_power_of_2(head_dim), + BLOCK_P=triton.next_power_of_2(max(max_rows, 1)), + BLOCK_C=triton.next_power_of_2(ngram_size - 1), + HAS_CARRY=has_initial_states is not None, + num_warps=4, + ) + return out + + def reference( + self, + input_ids: torch.Tensor, + context_pool: torch.Tensor, + state_indices: torch.Tensor, + has_initial_states: Optional[torch.Tensor], + query_start_loc: torch.Tensor, + table: torch.Tensor, + scale: torch.Tensor, + layer_multipliers: torch.Tensor, + head_vocab_sizes: torch.Tensor, + head_offsets: torch.Tensor, + eos_token_id: int, + max_rows: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """`(out, context_pool)` as the call leaves them, from the scalar + references of the two entries this replaces -- `ngram_token_history` + (window and published pool) and `ngram_hash` (rows) -- and a float64 + gather-and-scale rounded once. Mutates nothing.""" + carries = ( + torch.ones_like(state_indices, dtype=torch.bool) + if has_initial_states is None + else has_initial_states + ) + ngram_size = layer_multipliers.numel() + window, pool = ngram_token_history.reference( + input_ids, + context_pool, + state_indices, + carries, + query_start_loc, + eos_token_id, + ngram_size, + ) + rows = ngram_hash.reference( + window, layer_multipliers, head_vocab_sizes, head_offsets, eos_token_id, 1 + ).reshape(window.shape[0], -1) + gathered = table[rows.reshape(-1)].double().reshape(rows.shape[0], -1) + out = (gathered * scale.double().reshape(())).to(scale.dtype) + return out, pool + + def is_valid( + self, + input_ids: torch.Tensor, + context_pool: torch.Tensor, + state_indices: torch.Tensor, + has_initial_states: Optional[torch.Tensor], + query_start_loc: torch.Tensor, + table: torch.Tensor, + scale: torch.Tensor, + layer_multipliers: torch.Tensor, + head_vocab_sizes: torch.Tensor, + head_offsets: torch.Tensor, + eos_token_id: int, + max_rows: int, + ) -> None: + num_seqs = state_indices.shape[0] + ngram_size = layer_multipliers.numel() + assert input_ids.dim() == 1 and input_ids.is_contiguous(), "input_ids [T], contiguous" + assert ngram_size >= 2 and layer_multipliers.dtype == torch.int64, ( + "layer_multipliers int64, at least two orders" + ) + assert context_pool.dtype == torch.int64 and context_pool.stride(1) == 1, ( + "context_pool int64 with unit column stride" + ) + assert context_pool.shape[1] == ngram_size - 1, ( + f"context_pool must be ngram_size - 1 = {ngram_size - 1} wide" + ) + assert query_start_loc.shape[0] == num_seqs + 1, ( + "query_start_loc must be [num_seqs + 1]; a short one misattributes rows" + ) + assert has_initial_states is None or has_initial_states.shape[0] == num_seqs, ( + "has_initial_states must be one flag per sequence" + ) + assert head_vocab_sizes.numel() % (ngram_size - 1) == 0, ( + "heads must split evenly over the n-gram orders" + ) + assert head_offsets.numel() == head_vocab_sizes.numel(), "one offset per head" + assert head_vocab_sizes.dtype == head_offsets.dtype == torch.int64, "int64 vocab/offsets" + assert table.dim() == 2 and table.stride(1) == 1, "table [rows, D], unit column stride" + assert scale.numel() == 1, "one per-tensor scale" + assert max_rows >= 1, "max_rows >= 1" + + +ple_ngram_embed = PleNgramEmbed() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/bf16_to_fp32.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/bf16_to_fp32.py new file mode 100644 index 000000000000..d34139053140 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/bf16_to_fp32.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Widening bf16 -> fp32 of a contiguous tensor in one vectorized kernel: the logits +processor's cast of the step's `[rows, vocab]` logits. + +Tier 5 (Triton written for this target). torch's `.float()` runs this as an unrolled, +unvectorized `direct_copy` at under half the bandwidth; the values are the same bit for +bit. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: This checkpoint's vocabulary: the width of the logits rows. +VOCAB = 248320 +#: Elements a program (8 warps: 16 a thread, two 16-byte loads). +_BLOCK = 4096 +_WARPS = 8 + + +@triton.jit +def _bf16_to_fp32_kernel(x_ptr, y_ptr, n, BLOCK: tl.constexpr): + offs = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + mask = offs < n + tl.store(y_ptr + offs, tl.load(x_ptr + offs, mask=mask).to(tl.float32), mask=mask) + + +class Bf16ToFp32(OpWrapper): + """`x.float()` for a contiguous bf16 `x`: a fresh contiguous fp32 tensor of `x`'s + shape. Widening is exact, so every element is `x`'s value; NaN payloads widen as + torch widens them (the 16 bits shifted up). + + Fusion boundary: the cast only; the head's GEMM before it and the sampler after it + are the caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="decode logits: 128 rows of the 248320 vocabulary", spec=dict(shape=(128, VOCAB))), + Cell(why="mixed step logits: 124 rows", spec=dict(shape=(124, VOCAB))), + Cell(why="one row", spec=dict(shape=(1, VOCAB))), + Cell( + why="every bf16 bit pattern: NaNs, infinities, subnormals, signed zeros", + spec=dict(shape=(1, 65536), source="all_codes"), + ), + Cell(why="3 x 7: below one program", spec=dict(shape=(3, 7))), + Cell(why="rank 3, a ragged last program", spec=dict(shape=(2, 5, 4099))), + Cell(why="zero rows", spec=dict(shape=(0, VOCAB))), + ) + + note = """ + Metadata consumed: none; stateless. One Triton kernel, `cdiv(numel, 4096)` programs + of 8 warps over the flat tensor. Call site: the target's logits processor + (`_WideningLogitsProcessor`), on the head's bf16 logits; the shell's `.float()` + after it is then a no-op. + + GB300, graph replay: 128 x 248320 31.3 us against 82.1 for `.float()`, 124 rows + 30.7 against 79.9, 33 rows 8.3 against 16.5, one row 4.1 against 4.2. + + Preconditions: bf16, contiguous (the kernel reads the flat storage). + """ + + def __call__(self, x: torch.Tensor) -> torch.Tensor: + y = torch.empty(x.shape, dtype=torch.float32, device=x.device) + n = x.numel() + if n: + _bf16_to_fp32_kernel[(triton.cdiv(n, _BLOCK),)](x, y, n, BLOCK=_BLOCK, num_warps=_WARPS) + return y + + def reference(self, x: torch.Tensor) -> torch.Tensor: + return x.float() + + def is_valid(self, x: torch.Tensor) -> None: + assert x.dtype == torch.bfloat16, "bf16 input" + assert x.is_contiguous(), "contiguous input: the kernel reads the flat storage" + + +bf16_to_fp32 = Bf16ToFp32() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/fp8_quantize_1x128.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/fp8_quantize_1x128.py index b006f4e73adc..8129463ba804 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/fp8_quantize_1x128.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/fp8_quantize_1x128.py @@ -87,9 +87,13 @@ class Fp8Quantize1x128(OpWrapper): 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", + why="8192 tokens: max_num_tokens of the identity config", spec=dict(tokens=8192, **_CKPT), ), + Cell( + why="16384 rows: max_num_tokens of the tuned config", + spec=dict(tokens=16384, **_CKPT), + ), Cell( why="H 128: one whole block per row", spec=dict(tokens=7, hidden=128, use_ue8m0=False) ), @@ -138,7 +142,8 @@ class Fp8Quantize1x128(OpWrapper): independent; the registered fake matches eager shapes. Hard limits - T certified 0..8192 at H 2560; above that not probed + T certified 0..16384 at H 2560 (max_num_tokens of the tuned config); + above that not probed H multiple of 16 (raises otherwise); certified 16..7168 """ diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py index 973a12dc0501..0bd6a6662bbe 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py @@ -58,6 +58,11 @@ def region(x): gate = gate_buffer[:, :, d:] elif gate_layout == "packed": gate_buffer = gate = randn(t, heads, d) + elif gate_layout == "fused_row": + # The target's fused input projection row [q | k | v | gate | index | pad]: + # 14080 columns, the gate a 2-D [T, hidden] window at column 7168. + gate_buffer = randn(t, 14080) + gate = gate_buffer[:, 7168 : 7168 + hidden] else: gate_buffer = gate = randn(t, hidden, dt=spec.get("gate_dtype", dtype)) return buffer, region(buffer), region, gate_buffer, gate diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_silu_and_mul_rows.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_silu_and_mul_rows.py new file mode 100644 index 000000000000..7502b3963747 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_silu_and_mul_rows.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the silu_and_mul_rows entry: the cells bitwise `flashinfer_silu_and_mul` on a +contiguous copy of the gate|up columns (and so inside its reference band), the extra columns +unread, graph capture.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.activation.flashinfer_silu_and_mul import ( # noqa: E501 + flashinfer_silu_and_mul, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.activation.silu_and_mul_rows import ( # noqa: E501 + silu_and_mul_rows as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "silu_and_mul_rows requires a CUDA device" + +INTER = 640 + + +def _build(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + rows, width = spec["rows"], spec.get("width", 2 * INTER + 8) + x = (torch.randn(rows, width, generator=g, device="cuda") * 3.0).to(torch.bfloat16) + if rows: + x[0, :4] = torch.tensor([0.0, -0.0, 88.0, -100.0], device="cuda").to(torch.bfloat16) + return x + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + certifying(op) + x = _build(cell.spec, zlib.crc32(cell.why.encode())) + out = op(x, INTER) + assert out.shape == (x.shape[0], INTER) and out.is_contiguous() + torch.testing.assert_close(out, op.reference(x, INTER)) + if x.shape[0]: + assert torch.equal(out, flashinfer_silu_and_mul(x[:, : 2 * INTER].contiguous())) + + +def test_extra_columns_are_not_read() -> None: + x = _build(dict(rows=64), 7) + a = op(x, INTER) + x[:, 2 * INTER :] = float("nan") + assert torch.equal(op(x, INTER), a) + + +def test_cuda_graph_capture_and_replay() -> None: + x = _build(dict(rows=128), 8) + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + op(x, INTER) + torch.cuda.current_stream().wait_stream(warmup) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = op(x, INTER) + x.copy_(_build(dict(rows=128), 9)) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, op(x, INTER)) + + +def test_guards() -> None: + x = _build(dict(rows=4), 5) + with validating(op), pytest.raises(AssertionError): + op(x[:, : 2 * INTER - 2], INTER) + with validating(op), pytest.raises(AssertionError): + op(x.float(), INTER) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_bf16_gemv_pdl.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_bf16_gemv_pdl.py new file mode 100644 index 000000000000..0146f9a12052 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_bf16_gemv_pdl.py @@ -0,0 +1,129 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the bf16_gemv_pdl entry. + +Cells against the fp32 reference within the band `low_m_bf16_mm` uses (the +same products, summed in another order). Beyond the cells: a PDL chain under +CUDA-graph capture and replay, where each GEMV consumes the previous one's +output and so would read a stale activation if it read before its wait; PDL on +and off agree bitwise; the band rejects the near misses of a wrong layout; 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.gemm.bf16_gemv_pdl import ( # noqa: E501 + MAX_ROWS, + launch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.gemm.bf16_gemv_pdl import ( # noqa: E501 + bf16_gemv_pdl as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import ( # noqa: E402,F401 — path above; fixture by name + ULP, + assert_within_ulp, + certifying, + validating, +) + +assert torch.cuda.is_available(), "bf16_gemv_pdl requires a CUDA device" + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Accumulation order only: 4 bf16 ulp element-wise against the row scale, + 2 ulp of relative RMS, as for `low_m_bf16_mm` and `cublas_mm`.""" + assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0, ulp=ULP[torch.bfloat16]) + + +def _build(m, n, k, seed): + """`mat_b` is a `[N, K]` weight's `.t()`, as the target stores its projections.""" + g = torch.Generator(device="cuda").manual_seed(seed) + mat_a = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + weight = (torch.randn(n, k, generator=g, device="cuda") * k**-0.5).to(torch.bfloat16) + return mat_a, weight.t() + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 + certifying(op) + s = cell.spec + mat_a, mat_b = _build(s["m"], s["n"], s["k"], seed=zlib.crc32(cell.why.encode()) % 2**31) + out = op(mat_a, mat_b, s["pdl"]) + assert out.shape == (s["m"], s["n"]) and out.dtype == torch.bfloat16 and out.is_contiguous() + _compare(out, op.reference(mat_a, mat_b, s["pdl"])) + + +@pytest.mark.parametrize("pdl", [False, True]) +def test_dependent_chain_under_graph_replay(pdl) -> None: + """Decode runs this under CUDA graphs, back to back with other kernels. Each + link of the chain reads the output of the one before it: a GEMV that read + its activation before `griddepcontrol.wait` would see the previous replay's + values, and the comparison below would fail on the second replay.""" + dims = [2560, 6144, 2560, 640, 2560] # K of each link; N is the next link's K + weights = [_build(1, n, k, seed=i)[1] for i, (k, n) in enumerate(zip(dims, dims[1:] + [16480]))] + x = torch.randn(1, dims[0], device="cuda").to(torch.bfloat16) + + def chain(): + h = x + outs = [] + for w in weights: + h = op(h, w, pdl) + outs.append(h) + return outs + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(2): + chain() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outs = chain() + for trial in range(3): + g = torch.Generator(device="cuda").manual_seed(100 + trial) + x.copy_(torch.randn(x.shape, generator=g, device="cuda").to(x.dtype)) + graph.replay() + torch.cuda.synchronize() + h = x + for w, out in zip(weights, outs): + # Each link against the reference of what the previous link actually stored. + _compare(out, op.reference(h, w)) + h = out + + +@pytest.mark.parametrize("m,n,k", [(1, 16480, 2560), (2, 2560, 6144), (1, 2560, 640)]) +def test_pdl_on_and_off_agree_bitwise(m, n, k) -> None: + """PDL changes when the kernel may start, not what it computes.""" + mat_a, mat_b = _build(m, n, k, seed=11) + assert torch.equal(launch(mat_a, mat_b, True), launch(mat_a, mat_b, False)) + + +@pytest.mark.parametrize("m,n,k", [(1, 16480, 2560), (2, 2560, 6144)]) +def test_band_discriminates_wrong_variants(m, n, k) -> None: + """The band admits the kernel and rejects the near misses a wrong k-tile + split or layout would produce: half of K missing, a one-column shift of the + output, a 3% scale error.""" + mat_a, mat_b = _build(m, n, k, seed=7) + out = op(mat_a, mat_b) + ref = op.reference(mat_a, mat_b) + _compare(out, ref) + half = mat_a.clone() + half[:, k // 2 :] = 0 + for wrong in (op.reference(half, mat_b), torch.roll(ref, 1, dims=1), ref * 1.03): + with pytest.raises(AssertionError): + _compare(out, wrong) + + +def test_guard_refuses_a_row_major_mat_b_and_too_many_rows() -> None: + mat_a, mat_b = _build(1, 512, 2560, seed=3) + with validating(op), pytest.raises(AssertionError): + op(mat_a, mat_b.contiguous()) + a, b = _build(MAX_ROWS + 1, 512, 2560, seed=4) + with validating(op), pytest.raises(AssertionError): + op(a, b) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_block_scale_moe_renorm_permuted.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_block_scale_moe_renorm_permuted.py new file mode 100644 index 000000000000..fbdc5c576c43 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_block_scale_moe_renorm_permuted.py @@ -0,0 +1,118 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for ops/moe/fp4_block_scale_moe_renorm_permuted. + +Per cell: the emulated finalize of the three outputs is **bitwise** the +finalized entry's output on identical operands (so the fold a consumer makes +with the same arithmetic is bitwise what the runner's finalize stored), the +outputs have the documented shapes, dtypes and index map, and the finalized +rows sit inside the pre-routed entry's band against the fp32 reference. The +operand construction is the finalized entry's test's, loaded by path. +""" + +import importlib.util +import zlib +from pathlib import Path + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp4_block_scale_moe_renorm_permuted import ( # noqa: E501 + finalize, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp4_block_scale_moe_renorm_permuted import ( # noqa: E501 + fp4_block_scale_moe_renorm_permuted as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp4_block_scale_moe_renorm_routed import ( # noqa: E501 + fp4_block_scale_moe_renorm_routed as finalized, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "fp4_block_scale_moe_renorm_permuted requires a CUDA device" + + +def _routed_test(): + path = Path(__file__).resolve().parent / "test_modeling_v2_fp4_block_scale_moe_renorm_routed.py" + spec = importlib.util.spec_from_file_location("_fp4_moe_renorm_routed_test", path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +R = _routed_test() +H = R.H + + +def _args(args, data, sf, logits, scalars, spec, top_k=None): + return ( + logits, + data, + sf.view(torch.float8_e4m3fn), + args["gemm1_weights"], + args["gemm1_weights_scale"], + args["gemm2_weights"], + args["gemm2_weights_scale"], + *scalars, + spec["num_experts"], + spec["top_k"] if top_k is None else top_k, + spec["intermediate"], + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + spec = cell.spec + t, k, hidden = spec["tokens"], spec["top_k"], spec["hidden"] + seed = zlib.crc32(cell.why.encode()) % 2**31 + args, ref, data, sf, x, logits, ids, wts, g2, scalars = R._operands(spec, seed) + call = _args(args, data, sf, logits, scalars, spec) + + rows, topk_w, topk_rows = op(*call) + assert rows.dtype == torch.bfloat16 and rows.dim() == 2 and rows.shape[1] == hidden + assert topk_w.shape == (t, k) and topk_w.dtype == torch.bfloat16 + assert topk_rows.shape == (t, k) and topk_rows.dtype == torch.int32 + # Every slot maps to its own row (no expert parallelism: no -1 slots). + assert int(topk_rows.min()) >= 0 and int(topk_rows.max()) < rows.shape[0] + assert topk_rows.flatten().unique().numel() == t * k, "two slots share a permuted row" + # The weights are the routing's (as a set per token; the slot order is the runner's). + assert torch.equal(topk_w.sort(-1).values, wts.sort(-1).values), "routing weights differ" + + # 1. Bitwise the finalized entry once finalized with the runner's arithmetic. + y = finalize(rows, topk_w, topk_rows) + y_fin = finalized(*call) + differing = (y.view(torch.int16) != y_fin.view(torch.int16)).any(dim=1).sum().item() + assert differing == 0, f"{differing} of {t} rows differ from the finalized entry" + + # 2. The fp32 reference, through the pre-routed entry's gate. + exp = op.reference(logits, x, k, ref["up"], ref["gate"], ref["down"], g2) + H._compare(y, exp.to(torch.bfloat16)) + + +def test_uncertified_pair_refused(): + spec = op.CELLS[0].spec + args, _, data, sf, _, logits, _, _, _, scalars = R._operands(spec, 13) + with validating(op), pytest.raises(AssertionError, match="certified Renormalize"): + op(*_args(args, data, sf, logits, scalars, spec, top_k=9)) + + +def test_cuda_graph_capture_and_replay(): + spec = next(c.spec for c in op.CELLS if c.spec["tokens"] == 1) + args, _, data, sf, _, logits, _, _, _, scalars = R._operands(spec, 43) + call = list(_args(args, data, sf, logits, scalars, spec)) + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + for _ in range(3): + op(*call) + torch.cuda.current_stream().wait_stream(warmup) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + g_rows, g_w, g_idx = op(*call) + logits.copy_(torch.randn_like(logits.float()).to(logits.dtype)) + graph.replay() + torch.cuda.synchronize() + want = finalized(*call) + assert torch.equal(finalize(g_rows, g_w, g_idx), want) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py new file mode 100644 index 000000000000..971ebb067b05 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py @@ -0,0 +1,152 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for ops/moe/fp4_block_scale_moe_renorm_routed. + +Two claims per cell: the logits-routed call is **bitwise** the composition it +replaces (`renorm_moe_routing_op` -> pre-routed `fp4_block_scale_moe_runner`), +and it sits inside the pre-routed entry's band against this entry's own +fp32 reference. The NVFP4 operand construction (layout, activations, scalars, +the band) is the pre-routed entry's test's, loaded by path so the two tests +cannot drift apart on what a valid operand is. +""" + +import importlib.util +import zlib +from pathlib import Path + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp4_block_scale_moe_renorm_routed import ( # noqa: E501 + CERTIFIED_PAIRS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp4_block_scale_moe_renorm_routed import ( # noqa: E501 + fp4_block_scale_moe_renorm_routed as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.renorm_moe_routing_op import ( # noqa: E501 + renorm_moe_routing_op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "fp4_block_scale_moe_renorm_routed requires a CUDA device" + +DEV = "cuda" +G1 = 137.0 # activation global scale, as the pre-routed entry's test draws it + + +def _helpers(): + path = ( + Path(__file__).resolve().parents[4] + / "moe" + / "test_modeling_v2_fp4_block_scale_moe_runner.py" + ) + spec = importlib.util.spec_from_file_location("_fp4_moe_runner_test_helpers", path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +H = _helpers() + + +def _logits(spec, gen): + t, e = spec["tokens"], spec["num_experts"] + if spec["dist"] == "coarse": + v = torch.randint(-1, 2, (t, e), generator=gen, device=DEV).float() + else: + v = torch.randn(t, e, generator=gen, device=DEV) + return v.to(spec["logits_dtype"]).contiguous() + + +def _operands(spec, seed): + gen = torch.Generator(device=DEV).manual_seed(seed) + args, ref = H._many_experts() + data, sf, x = H._rand_activation(spec["tokens"], spec["hidden"], G1, gen) + logits = _logits(spec, gen) + ids, wts = renorm_moe_routing_op(logits, spec["top_k"], torch.bfloat16) + g2 = H._calibrate_g2(x, ids, ref) + scalars = H._scalars(spec["num_experts"], G1, g2) + return args, ref, data, sf, x, logits, ids, wts, g2, scalars + + +def _routed(args, data, sf, logits, scalars, spec, top_k=None): + return op( + logits, + data, + sf.view(torch.float8_e4m3fn), + args["gemm1_weights"], + args["gemm1_weights_scale"], + args["gemm2_weights"], + args["gemm2_weights_scale"], + *scalars, + spec["num_experts"], + spec["top_k"] if top_k is None else top_k, + spec["intermediate"], + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + spec = cell.spec + seed = zlib.crc32(cell.why.encode()) % 2**31 + args, ref, data, sf, x, logits, ids, wts, g2, (s1, sg, s2) = _operands(spec, seed) + + y = _routed(args, data, sf, logits, (s1, sg, s2), spec) + assert y.shape == (spec["tokens"], spec["hidden"]) and y.dtype == torch.bfloat16 + + # 1. Bitwise the composition it replaces, on identical operands. + y_pre = H._call( + data, + sf, + args, + spec["num_experts"], + spec["top_k"], + topk_ids=ids, + topk_weights=wts, + output1_scale_scalar=s1, + output1_scale_gate_scalar=sg, + output2_scale_scalar=s2, + )[0] + differing = (y.view(torch.int16) != y_pre.view(torch.int16)).any(dim=1).sum().item() + assert differing == 0, f"{differing} of {spec['tokens']} rows differ from the composition" + + # 2. Its own reference, through the pre-routed entry's gate. + exp = op.reference(logits, x, spec["top_k"], ref["up"], ref["gate"], ref["down"], g2) + H._compare(y, exp.to(torch.bfloat16)) + + +def test_coarse_cell_exercises_tie_breaking(): + """The coarse cell is only worth its slot if ties decide real slots: the + opposite tie order (ties to the HIGHER expert index) must pick a different + expert set on most rows, so the cell's bitwise check pins the direction.""" + spec = next(c.spec for c in op.CELLS if c.spec["dist"] == "coarse") + e, k = spec["num_experts"], spec["top_k"] + gen = torch.Generator(device=DEV).manual_seed(7) + logits = _logits(spec, gen) + ids, _ = renorm_moe_routing_op(logits, k, torch.bfloat16) + flipped = torch.flip(logits, [-1]) + order = torch.sort(flipped, dim=-1, descending=True, stable=True).indices[:, :k] + high = (e - 1 - order).int() + differ = (ids.sort(-1).values != high.sort(-1).values).any(-1).float().mean().item() + assert differ > 0.5, f"only {differ:.2f} of rows break ties differently" + + +def test_strided_logits_refused(): + # More than one row: a one-row column slice is contiguous by definition. + spec = next(c.spec for c in op.CELLS if c.spec["tokens"] == 8) + args, _, data, sf, _, logits, _, _, _, scalars = _operands(spec, 11) + wide = torch.zeros(spec["tokens"], spec["num_experts"] + 8, device=DEV, dtype=logits.dtype) + wide[:, : spec["num_experts"]] = logits + with validating(op), pytest.raises(AssertionError, match="routing_logits must be contiguous"): + _routed(args, data, sf, wide[:, : spec["num_experts"]], scalars, spec) + + +def test_uncertified_pair_refused(): + spec = op.CELLS[0].spec + assert (spec["num_experts"], 9) not in CERTIFIED_PAIRS + args, _, data, sf, _, logits, _, _, _, scalars = _operands(spec, 13) + with validating(op), pytest.raises(AssertionError, match="certified Renormalize"): + _routed(args, data, sf, logits, scalars, spec, top_k=9) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_gate_conv_inject_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_gate_conv_inject_norm.py new file mode 100644 index 000000000000..07a42aea5272 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_gate_conv_inject_norm.py @@ -0,0 +1,298 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the ple_gate_conv_inject_norm entry. + +Cells are gated against the float64 reference that rounds where the op rounds: +the stream (`hc_out`) and the published conv state within one ulp of each +element, the normed stream within two (see `_compare`). Beyond the cells: +agreement with the composition it replaces at the same band, separation from +plausible wrong formulas, 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_hc_inject_norm import ( # noqa: E501 + qwen4_exp_hc_inject_norm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.norm.flashinfer_gemma_rmsnorm import ( # noqa: E501 + flashinfer_gemma_rmsnorm, +) +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, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.ple_gate_conv_inject_norm import ( # noqa: E501 + CKPT_CLAMP_MIN, + CKPT_DILATION, + CKPT_HC_COUNT, + CKPT_HIDDEN, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.ple_gate_conv_inject_norm import ( # noqa: E501 + ple_gate_conv_inject_norm as op, +) +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, +) + +__extra_import_path__ = ["../../../.."] +from _validating import ( # noqa: E402,F401 — path above; fixture by name + assert_within_ulp, + certifying, + validating, +) + +assert torch.cuda.is_available(), "ple_gate_conv_inject_norm requires a CUDA device" + +DEV = "cuda" +EPS = 1e-6 +PAD = -1 +GAMMA_RMS = 0.477 + + +def _ulp_ratio(out: torch.Tensor, ref: torch.Tensor) -> torch.Tensor: + """Per-element distance from `ref` in units of one ulp of `|ref|`.""" + eps = torch.finfo(ref.dtype).eps + tiny = torch.finfo(ref.dtype).tiny + ref64 = ref.double() + return (out.double() - ref64).abs() / (eps * ref64.abs() + tiny) + + +def _stats(out: torch.Tensor, ref: torch.Tensor) -> tuple[float, float]: + ratio = _ulp_ratio(out, ref) + return ratio.max().item(), (out != ref).double().mean().item() + + +#: The band. Every rounding point of the composition is kept, so the op and the +#: float64 reference differ only where an fp32 reduction (three sums of squares +#: and the gate's dot per stream) or fp32 `exp` lands a value across a bf16 +#: midpoint -- measured on 2e-5..2.3e-4 of the elements. A flipped upstream +#: element then propagates, and through `hc + ple` it meets cancellation +#: (`hc ~ -ple`), so a per-element ulp is unbounded there: 73 ulp of a +#: 0.003-sized element in a row whose largest is ~10. The stream and its norm +#: are therefore gated in units of their row's scale (`assert_within_ulp`): +#: 2 ulp of the row maximum per element and 0.1 ulp relative RMS. The +#: composition this replaces sits at 0.87 / 0.33 of a row-max ulp against the +#: same reference (measured: normed / hc_out at 128 rows), the op at 0.72 / +#: 0.38. The published conv state is one rounding of the conv-input norm, +#: whose only perturbation is a one-ulp flip of a gated element: 2 ulp per +#: element (measured worst 1.78, the old composition 1.56). +ROW_ULP, RMS_ULP, STATE_ULP = 2.0, 0.1, 2.0 +#: The per-element row-max gate is a tail statistic, and a context-shaped cell +#: draws ~100x more elements than a 128-row one. Against the same float64 +#: reference over three seeds at 16384 rows, the composition this replaces +#: reaches 2.91 row-max ulp (normed) and the op 2.94; at 8192 rows 1.89 / 1.89 +#: (`microbench/diag_band.py` in the candidate dir). Cells above 1024 rows are +#: therefore gated at 4; the RMS gate is unchanged. +CONTEXT_ROWS, CONTEXT_ROW_ULP = 1024, 4.0 + + +def _compare(out, ref) -> None: + row_ulp = CONTEXT_ROW_ULP if out[0].shape[0] > CONTEXT_ROWS else ROW_ULP + for name, got, want in zip(("hc_out", "normed", "conv_states"), out, ref): + assert got.dtype == want.dtype and got.shape == want.shape, name + if not got.numel(): + continue + if name == "conv_states": + worst = _ulp_ratio(got, want).max().item() + assert worst <= STATE_ULP, f"{name}: {worst:.4g} ulp" + else: + assert_within_ulp(got, want, row_ulp, RMS_ULP) + + +def _build(spec, seed): + g = torch.Generator(device=DEV).manual_seed(seed) + gc = torch.Generator().manual_seed(seed) + lengths = spec["lengths"] + c = spec.get("hc_count", CKPT_HC_COUNT) + d = spec.get("hidden", CKPT_HIDDEN) + dtype = spec.get("dtype", torch.bfloat16) + w = c * d + rows = sum(lengths) + num_seqs = len(lengths) + pool_slots = spec.get("pool_slots", num_seqs + 2) + + def rand(*shape, scale=1.0): + return (torch.randn(*shape, generator=g, device=DEV) * scale).to(dtype) + + # Streams at different scales, as the trained residual's are. + scales = torch.tensor([0.5 * 2.0 ** (i % 4) for i in range(c)], device=DEV) + hc = (rand(rows, w, scale=3.0).float() * scales.repeat_interleave(d)).to(dtype) + key = rand(rows, w) + if spec.get("layout") == "row_slice": + hbuf, kbuf = rand(2 * rows, w), rand(2 * rows, w) + hbuf[::2], kbuf[::2] = hc, key + hc, key = hbuf[::2], kbuf[::2] + value = rand(rows, d, scale=0.5) + norms = [rand(w, scale=GAMMA_RMS) for _ in range(4)] + conv_weight = rand(w, 1, 4, scale=0.3) + conv_states = rand(pool_slots, w, 3 * CKPT_DILATION) + order = (torch.randperm(pool_slots - 2, generator=gc) + 1).tolist() + slots = ([0, pool_slots - 1] + order)[:num_seqs] if pool_slots >= 64 else list(range(num_seqs)) + if "pad" in spec: + slots[spec["pad"]] = PAD + slots = torch.tensor(slots, dtype=torch.int32, device=DEV) + starts = [0] + for length in lengths: + starts.append(starts[-1] + length) + qsl = torch.tensor(starts, dtype=torch.int64, device=DEV) + max_rows = spec.get("max_rows", max(lengths)) + carries = spec.get("carries") + carries = None if carries is None else torch.tensor(carries, dtype=torch.bool, device=DEV) + return [ + hc, key, value, norms[0], norms[1], norms[2], conv_weight, conv_states, slots, qsl, + norms[3], EPS, c, CKPT_CLAMP_MIN, CKPT_DILATION, PAD, max_rows, carries, + ] # fmt: skip + + +def _run(args): + """`(hc_out, normed, conv_states)` from the op, on a copy of the pool.""" + args = list(args) + args[7] = args[7].clone() + hc_out, normed = op(*args) + return hc_out, normed, args[7] + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) + hc_before = args[0].clone() + ref = op.reference(*args) + out = _run(args) + assert torch.equal(args[0], hc_before), "hc is read-only" + assert out[0].is_contiguous() and out[1].is_contiguous() + _compare(out, ref) + for name, got, want in zip(("hc_out", "normed", "conv_states"), out, ref): + worst, frac = _stats(got, want) + print( + f"\n {cell.why[:40]:<40} {name:<12} worst {worst:.3g} ulp, differ {frac:.2e}", end="" + ) + + +def _grouped(x, w, c): + d = x.shape[1] // c + parts = torch.split(x, [d] * c, 1) + return torch.cat( + [flashinfer_gemma_rmsnorm(p, s, EPS) for s, p in zip(torch.split(w, d, 0), parts)], 1 + ) + + +def _old( + hc, + key, + value, + nk, + nq, + nc, + cw, + states, + slots, + qsl, + nx, + eps, + c, + clamp, + dil, + pad, + _p, + carries=None, +): + """The calls the decode / prefill PLE block and the site after it made before this entry.""" + gated = gated_ngram_injection(_grouped(key, nk, c), _grouped(hc, nq, c), value, c, clamp) + convolved = dilated_short_conv_paged( + x=torch.transpose(_grouped(gated, nc, c), 0, 1), + weight=cw, + bias=None, + conv_states=states, + query_start_loc=qsl, + cache_indices=slots, + has_initial_state=torch.ones_like(slots, dtype=torch.bool) if carries is None else carries, + dilation=dil, + silu_activation=True, + pad_slot_id=pad, + ) + hc_new = torch.add(hc, torch.add(gated, torch.transpose(convolved, 0, 1))) + _, normed = qwen4_exp_hc_inject_norm(hc_new, None, None, nx, eps, c) + return hc_new, normed, states + + +@pytest.mark.parametrize( + "lengths,carries", + [ + ((1,), None), + ((1,) * 32, None), + ((1,) * 128, None), + ((4,) * 8, None), + ((8192,), (False,)), + ((5000, 3000) + (1,) * 32, (False, True) + (True,) * 32), + ], + ids=["m1", "m32", "m128", "mtp4x8", "prefill8192", "mixed2ctx32dec"], +) +def test_matches_the_composition_it_replaces(lengths, carries) -> None: + args = _build(dict(lengths=lengths, pool_slots=461, carries=carries), seed=len(lengths)) + old_args = list(args) + old_args[7] = args[7].clone() + want = _old(*old_args) + got = _run(args) + _compare(got, want) + for name, a, b in zip(("hc_out", "normed", "conv_states"), got, want): + worst, frac = _stats(a, b) + print( + f"\n {str(len(lengths)):>4} seqs {name:<12} worst {worst:.3g} ulp, differ {frac:.2e}", + end="", + ) + + +def test_wrong_variants_land_outside_the_band() -> None: + """Each mis-reading keeps every shape; each lands far outside the band.""" + args = _build(dict(lengths=(1,) * 16, pool_slots=461), seed=21) + want = op.reference(*args) + variants = { + "dilation 1": dict(index=14, value=1), + "clamp floor 1.0": dict(index=13, value=1.0), + "next norm = query norm": dict(index=10, value=args[4]), + "conv norm = key norm": dict(index=5, value=args[3]), + } + for name, change in variants.items(): + wrong = list(args) + wrong[change["index"]] = change["value"] + if name == "dilation 1": + wrong[7] = wrong[7][:, :, :3].contiguous() + got = _run(wrong) + worst = max(_ulp_ratio(got[i], want[i]).max().item() for i in (0, 1)) + print(f"\n {name:<26} {worst:10.4g} ulp", end="") + assert worst > 5.0, name + + +def test_guard_refuses_a_short_value() -> None: + args = _build(dict(lengths=(1, 1)), seed=4) + args[2] = args[2][:, :-1] + with validating(op), pytest.raises(AssertionError): + op(*args) + + +def test_cuda_graph_capture_and_replay() -> None: + args = _build(dict(lengths=(1,) * 8, pool_slots=461), seed=8) + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + pool = args[7].clone() + with torch.cuda.stream(warmup): + for _ in range(2): + op(*args) + torch.cuda.current_stream().wait_stream(warmup) + args[7].copy_(pool) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_hc, graph_normed = op(*args) + fresh = _build(dict(lengths=(1,) * 8, pool_slots=461), seed=18) + args[0].copy_(fresh[0]) + args[1].copy_(fresh[1]) + args[7].copy_(pool) + eager = _run(args) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(graph_hc, eager[0]) + assert torch.equal(graph_normed, eager[1]) + assert torch.equal(args[7], eager[2]) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_bf16_to_fp32.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_bf16_to_fp32.py new file mode 100644 index 000000000000..7196d321e733 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_bf16_to_fp32.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the bf16_to_fp32 entry: the cells bit for bit `x.float()`, graph capture, the +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.quantization.bf16_to_fp32 import ( # noqa: E501 + bf16_to_fp32 as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "bf16_to_fp32 requires a CUDA device" + + +def _build(spec, seed): + shape = spec["shape"] + if spec.get("source") == "all_codes": + return ( + torch.arange(-32768, 32768, dtype=torch.int32, device="cuda") + .to(torch.int16) + .view(torch.bfloat16) + .reshape(shape) + ) + g = torch.Generator(device="cuda").manual_seed(seed) + return (torch.randn(shape, generator=g, device="cuda") * 8).to(torch.bfloat16) + + +def _bitwise(a: torch.Tensor, b: torch.Tensor) -> bool: + return a.shape == b.shape and torch.equal(a.view(torch.int32), b.view(torch.int32)) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + certifying(op) + x = _build(cell.spec, zlib.crc32(cell.why.encode())) + out = op(x) + ref = op.reference(x) + assert out.dtype == torch.float32 and out.is_contiguous() + assert _bitwise(out, ref), "widening is exact; a bit difference is a defect" + + +def test_cuda_graph_capture_and_replay() -> None: + x = _build(dict(shape=(128, 4099)), 1) + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + op(x) + torch.cuda.current_stream().wait_stream(warmup) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = op(x) + x.copy_(_build(dict(shape=(128, 4099)), 2)) + graph.replay() + torch.cuda.synchronize() + assert _bitwise(out, x.float()) + + +def test_guards() -> None: + x = _build(dict(shape=(8, 64)), 3) + with validating(op), pytest.raises(AssertionError): + op(x.float()) + with validating(op), pytest.raises(AssertionError): + op(x[:, ::2]) From 614965ce7f1fafc6339029fcd598f7fcb1b30a7d Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Sun, 18 Oct 2026 23:00:05 -0700 Subject: [PATCH 09/28] [TRTLLM-16460][perf] Qwen3.8-Flash-Next target: run the fast paths `modeling.py` and `weights.py` take the ops of the previous commits: the decode step as PDL chains (attention prologue, sparse decode split-K, GDN projection + conv + recurrence + gated norm, HC sites, the MoE site), the context step on the row-tile attention, the GDN prefill conv pass, the CuTe DSL MoE for prefill-sized launches and the shared-expert gate logit in the gate|up GEMM, the host-offloaded PLE table mapped from registered pageable memory, the cuBLAS HC first projection over 352 of the 384 padded columns, and the logits widened in one kernel. The fast paths run over a bf16 KV pool and an fp32 GDN state; FP8/NVFP4 KV and a bf16 state keep the feature commit's paths. With the tests that drive the target's modeling or weights. Folds these commits of the review branch (their changes to the files of this commit): - qwen3.8 target: carry the decode-fusion perf work onto the feature target (62151b727b) - qwen3.8 target: canonical decode selection in one launch (ac8ad93f0f) - qwen3.8 target: canonical context-row selection in one launch (bd29d1c90f) - qwen3.8 target: map the host-offloaded PLE table from registered pageable memory (0337d77eb7) - qwen3.8 target: stream the GDN decode step with its gated norm (15dfd83ba3) - qwen3.8 target: shared-expert gate logit from the gate|up GEMM at prefill (2ad541143d) - qwen3.8 target: a mixed step's generation rows on the streaming GDN kernel (73d22b699f) - qwen3.8 target: cuBLAS HC first projection over 352 of the 384 padded columns (8dfc993426) - qwen3.8 target: widen the logits to fp32 in one vectorized kernel (edb0d2e93a) Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../modeling.py | 2237 ++++++++++++----- .../weights.py | 283 ++- .../test_modeling_v2_cute_dsl_weight_path.py | 5 +- ...est_modeling_v2_fp4_moe_renorm_cute_dsl.py | 164 ++ .../ple/test_modeling_v2_ple_ngram_embed.py | 223 ++ .../ssm/test_modeling_v2_gdn_block.py | 407 +++ .../test_modeling_v2_logits_widening.py | 50 + ...test_modeling_v2_qwen4_exp_weight_table.py | 111 +- 8 files changed, 2715 insertions(+), 765 deletions(-) create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_ngram_embed.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_block.py create mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_logits_widening.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py index 16944580ba21..6496ccd141bb 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py @@ -20,6 +20,19 @@ * **Hyper-connections.** The inter-block residual is ``4 * 2560 = 10240`` wide. There is **no** ``model.language_model.norm.weight``: the stack-final ``hyper_connection_mixer`` is the final norm and ``lm_head`` reads its output. + In the trunk each site is ``ops/hc/qwen4_exp_hc_inject_norm`` (the previous block's + accumulate fused with this site's grouped norm; after a MoE block it also forms that + block's output from the routed and shared outputs and the gate logit) and then the + mixer: ``catalog/gemm/cublas_mm`` over the concatenated ``[down; block_inject]`` + weight, ``ops/hc/qwen4_exp_hc_lowrank_act`` and ``ops/hc/qwen4_exp_hc_up_gate_mean`` + (the up projection as a tcgen05 GEMM whose epilogue applies the sigmoid gate and the + stream mean). ``DecodeTarget`` mixers run the first projection as one tcgen05 launch + (``ops/hc/qwen4_exp_hc_decode_tc_first``, ``_HC_TC_FIRST_MIN_ROWS`` to + ``_HC_TC_FIRST_MAX_ROWS`` rows) or, at ``_HC_DECODE_GEMV_MAX_ROWS`` rows, two GEMV + launches (``ops/hc/qwen4_exp_hc_decode_first_cute``, ``ops/hc/qwen4_exp_hc_decode_mix``); + at the 48 MoE sites the mixer's last launch also writes the routed experts' NVFP4 + input. The MTP layer and the prefill PLE block keep the unfused + ``qwen4_exp_gated_residual`` / ``qwen4_exp_post_inject`` / ``_grouped_norm``. * **QSA sparse attention.** The index plane is maintained on every step so a request can cross ``dense_seq_len_threshold`` later without rebuilding its prefix; sparse selection replaces the dense core only above it. The MTP draft @@ -30,7 +43,11 @@ trunk's routed experts run on trtllm-gen's block-scale runner (``TRTLLM``, and what ``AUTO`` selects) or on the CuteDSL sort / grouped-GEMM pipeline (``CUTEDSL``), chosen by ``moe_config.backend``; any other request is refused. The MTP draft - layer's FP8 block-scale experts run on trtllm-gen under either. + layer's FP8 block-scale experts run on trtllm-gen under either. Under ``TRTLLM`` the + prefill-sized launches (``_CUTE_DSL_MOE_MIN_ROWS``) run the CuteDSL pipeline over a + second copy of the experts in its layout. A graph-captured decode step routes its block + projections through ``catalog/gemm/low_m_bf16_mm`` (FlashInfer low-M kernels at + M <= 32 on the shapes it lists, cuBLAS otherwise); see ``DecodeTarget.forward``. Integration facts this file depends on: @@ -107,20 +124,22 @@ thop_attention, ) from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.low_m_bf16_mm import low_m_bf16_mm from tensorrt_llm._torch._experimental.modeling_v2.catalog.graph.eager_on_graph import ( eager_on_graph, ) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fp4_block_scale_moe_runner import ( - fp4_block_scale_moe_runner, -) from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( fp4_quantize, ) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.flashinfer_chunk_gated_delta_rule import ( # noqa: E501 + flashinfer_chunk_gated_delta_rule, +) from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_utils import DecoderModel, DecoderModelForCausalLM from tensorrt_llm._torch.modules.fla.cached_replay import CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE +from tensorrt_llm._torch.modules.logits_processor import LogitsProcessor from tensorrt_llm._torch.speculative import get_spec_worker, should_use_separate_draft_kv_cache from tensorrt_llm._torch.utils import ActivationType from tensorrt_llm.logger import logger @@ -130,6 +149,7 @@ from .ops.activation.fused_sigmoid_mul_inplace import fused_sigmoid_mul_inplace from .ops.activation.gelu_erf import gelu_erf from .ops.activation.gelu_tanh import gelu_tanh +from .ops.activation.silu_and_mul_rows import silu_and_mul_rows from .ops.attention.qsa_canonical_selection import qsa_canonical_selection from .ops.attention.qsa_expand_block_indices import qsa_expand_block_indices from .ops.attention.qsa_mtp_reuse_topk import qsa_mtp_reuse_topk @@ -138,17 +158,41 @@ from .ops.attention.qsa_paged_kv_store import qsa_paged_kv_store from .ops.attention.qsa_paged_kv_store_nvfp4 import qsa_paged_kv_store_nvfp4 from .ops.attention.qsa_paged_sparse_gqa import qsa_paged_sparse_gqa +from .ops.attention.qsa_paged_sparse_gqa_decode import MAX_ROWS as _SPARSE_DECODE_MAX_ROWS +from .ops.attention.qsa_paged_sparse_gqa_decode import qsa_paged_sparse_gqa_decode +from .ops.attention.qsa_paged_sparse_gqa_decode import supports as _sparse_decode_supports from .ops.attention.qsa_paged_sparse_gqa_nvfp4 import qsa_paged_sparse_gqa_nvfp4 from .ops.attention.qsa_prefill_compress import qsa_prefill_compress +from .ops.attention.qsa_sparse_gqa_rowtile import qsa_sparse_gqa_rowtile +from .ops.attention.qsa_topk_expand_decode import qsa_topk_expand_decode from .ops.attention.qsa_topk_select import qsa_topk_select +from .ops.attention.qwen4_exp_attn_decode_prologue import qwen4_exp_attn_decode_prologue +from .ops.attention.qwen4_exp_qk_norm_rope import qwen4_exp_qk_norm_rope from .ops.attention.vision_rope_2d import vision_rope_2d from .ops.embedding.ngram_hash import ngram_hash +from .ops.gemm.bf16_gemv_pdl import TILING as _GEMV_TILED_K +from .ops.gemm.bf16_gemv_pdl import bf16_gemv_pdl from .ops.hc.qwen4_exp_gated_residual import qwen4_exp_gated_residual +from .ops.hc.qwen4_exp_hc_decode_first_cute import qwen4_exp_hc_decode_first_cute +from .ops.hc.qwen4_exp_hc_decode_mix import qwen4_exp_hc_decode_mix +from .ops.hc.qwen4_exp_hc_decode_mix_fp4_cute import ( + qwen4_exp_hc_decode_mix_cute, + qwen4_exp_hc_decode_mix_fp4_cute, +) +from .ops.hc.qwen4_exp_hc_decode_tc_first import qwen4_exp_hc_decode_tc_first +from .ops.hc.qwen4_exp_hc_gate_mean import qwen4_exp_hc_gate_mean +from .ops.hc.qwen4_exp_hc_inject_norm import qwen4_exp_hc_inject_norm +from .ops.hc.qwen4_exp_hc_lowrank_act import qwen4_exp_hc_lowrank_act +from .ops.hc.qwen4_exp_hc_up_gate_mean import qwen4_exp_hc_up_gate_mean +from .ops.hc.qwen4_exp_hc_up_gate_mean_fp4 import qwen4_exp_hc_up_gate_mean_fp4 from .ops.hc.qwen4_exp_post_inject import qwen4_exp_post_inject from .ops.moe.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion import ( cute_dsl_nvfp4_gather_grouped_gemm_act_fusion, ) from .ops.moe.cute_dsl_nvfp4_grouped_gemm import cute_dsl_nvfp4_grouped_gemm +from .ops.moe.fp4_block_scale_moe_renorm_permuted import fp4_block_scale_moe_renorm_permuted +from .ops.moe.fp4_block_scale_moe_renorm_routed import fp4_block_scale_moe_renorm_routed +from .ops.moe.fp4_moe_renorm_cute_dsl import fp4_moe_renorm_cute_dsl from .ops.moe.fp8_block_scale_moe_runner import fp8_block_scale_moe_runner from .ops.moe.fused_sigmoid_gate_mul_add import fused_sigmoid_gate_mul_add from .ops.moe.moe_sort import moe_sort @@ -159,6 +203,9 @@ from .ops.norm.rms_norm_gated_token_major import rms_norm_gated_token_major from .ops.ple.gated_ngram_injection import gated_ngram_injection from .ops.ple.ngram_token_history import ngram_token_history +from .ops.ple.ple_gate_conv_inject_norm import ple_gate_conv_inject_norm +from .ops.ple.ple_ngram_embed import ple_ngram_embed +from .ops.quantization.bf16_to_fp32 import bf16_to_fp32 from .ops.quantization.dequantize_e4m3_per_tensor import dequantize_e4m3_per_tensor from .ops.quantization.fp8_quantize_1x128 import fp8_quantize_1x128 from .ops.ssm.causal_conv1d_fwd import causal_conv1d_fwd @@ -175,6 +222,12 @@ ) from .ops.ssm.fused_recurrent_gated_delta_rule_update import fused_recurrent_gated_delta_rule_update from .ops.ssm.fused_sigmoid_gating_delta_rule_update import fused_sigmoid_gating_delta_rule_update +from .ops.ssm.gdn_decode_gated_norm_pdl import MAX_ROWS as _GDN_DECODE_FUSED_MAX_ROWS +from .ops.ssm.gdn_decode_gated_norm_pdl import gdn_decode_gated_norm_pdl +from .ops.ssm.gdn_decode_rows_gated_norm import gdn_decode_rows, gdn_decode_rows_gated_norm +from .ops.ssm.gdn_in_proj_conv1d_pdl import MAX_ROWS as _GDN_CONV_FUSED_MAX_ROWS +from .ops.ssm.gdn_in_proj_conv1d_pdl import gdn_in_proj_conv1d_pdl +from .ops.ssm.gdn_prefill_conv_post import gdn_prefill_conv_post from .ops.ssm.pack_gdn_decode_qkv import pack_gdn_decode_qkv from .ops.ssm.reset_recurrent_state_rows import reset_recurrent_state_rows from .ops.vision.patch_embed_conv3d import patch_embed_conv3d @@ -206,6 +259,33 @@ ) +# A context launch at or above this many rows takes `qsa_sparse_gqa_rowtile`; +# measured on GB300 it loses at 543 rows and wins from 4096 rows up. +_ROWTILE_MIN_ROWS = 2048 # capacity: fixed: a dispatch threshold, not a buffer size + +#: Context-step MoE launches of at least this many rows run on the CuTe DSL MoE +#: (`fp4_moe_renorm_cute_dsl`, from the second expert copy `weights.py` keeps in +#: its layout); smaller ones, every decode one among them, keep the trtllm-gen +#: runner. From GB300 micro-benches, us a layer, CuTe DSL vs trtllm-gen: 2048 +#: rows 275 vs 273, 8192 rows 451 vs 555, 16384 rows 656 vs 882. +_CUTE_DSL_MOE_MIN_ROWS = 4096 # capacity: fixed: a dispatch threshold, not a buffer size + + +def _decode_projection(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor: + """A graph-captured decode step's bf16 projection, always a PDL launch. + + One row runs `bf16_gemv_pdl` on every K it tiles (each decode projection of this + checkpoint): its weight streams into registers before it waits on the producer, and it + triggers its own dependents. More rows run `low_m_bf16_mm` with its FlashInfer + branches' PDL on (its cuBLAS branch, above M=32 and off its table, ignores the flag). + Every kernel either path launches reads the activation only after + `griddepcontrol.wait`. + """ + if mat_a.shape[0] == 1 and mat_b.shape[0] in _GEMV_TILED_K: + return bf16_gemv_pdl(mat_a, mat_b, True) + return low_m_bf16_mm(mat_a, mat_b, True) + + # Per-call constants of the dense attention call shape: packed-QKV causal GQA, RoPE # applied outside (fused_qk_norm_rope), bf16, fp8-e4m3 or NVFP4 pool. The sparse / skip-softmax # group stays None/0: QSA calls selection and sparse attention as separate kernels and @@ -405,7 +485,7 @@ _PAD_SLOT_ID = -1 # Gated-DeltaNet gating form: the only pair `fused_sigmoid_gating_delta_rule_update` -# is certified at; matches `fused_gdn_post_conv`'s default so prefill and decode +# is certified at; matches `gdn_prefill_conv_post`'s default so prefill and decode # compute the same `g`. _SOFTPLUS_BETA = 1.0 _SOFTPLUS_THRESHOLD = 20.0 @@ -463,6 +543,31 @@ def _ple_host_offload_enabled() -> bool: #: draft rejected; only the acceptance gate can tell. _MTP_EMBED_BLOCK_ALL_STREAMS = True +#: Decode batches up to this many rows run each trunk hyper-connection site as +#: `qwen4_exp_hc_decode_first_cute` + `qwen4_exp_hc_decode_mix` (two GEMV launches); +#: larger ones as `qwen4_exp_hc_inject_norm` + cuBLAS. A dispatch threshold from the +#: GB300 micro-bench (the GEMVs win at 1 row and lose from 2), at most the entries' +#: MAX_ROWS. +_HC_DECODE_GEMV_MAX_ROWS = 1 # capacity: fixed: dispatch threshold, not a buffer bound + +#: Decode row count from which the hyper-connection up projection runs as +#: `qwen4_exp_hc_up_gate_mean` (gate and stream mean in the GEMM epilogue). The +#: fused kernel's 80 tiles read the 6.5 MB up weight slower than cuBLAS at +#: one row (5.8 vs 5.2 us); from 2 rows on the saved launch and logits round +#: trip win (5.6 vs 6.1 us at 2, 5.5 vs 7.4 at 128; GB300, graph-replayed, +#: cold weights). Prefill always takes the fused op. Trunk sites at most +#: `_HC_DECODE_GEMV_MAX_ROWS` rows never reach the mixer (two GEMV launches). +_HC_FUSED_UP_MIN_ROWS = 2 # capacity: fixed: dispatch threshold, not a buffer bound + +#: Decode row counts whose hyper-connection mixer runs its first half (projection +#: and activations) as the one tcgen05 launch `qwen4_exp_hc_decode_tc_first` +#: instead of cuBLAS split-K (+ its reduce) + `qwen4_exp_hc_lowrank_act`. The upper +#: bound is the entry's MAX_ROWS (one 128-row MMA tile); larger batches keep cuBLAS. +_HC_TC_FIRST_MIN_ROWS = 2 # capacity: fixed: dispatch threshold, not a buffer bound +_HC_TC_FIRST_MAX_ROWS = ( + 128 # capacity: fixed: dispatch threshold (one MMA tile), not a buffer bound +) + #: FP8 block-scale block edge, on both weight-scale axes and the activation row. _FP8_BLOCK = _weights.FP8_BLOCK #: Values per e4m3 block scale in an NVFP4 KV pool. @@ -937,6 +1042,13 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig f"residual and is a different assembly" ) self.stream = self.hc_count * self.hidden + # Columns of the padded first-projection weight (`_fuse_hc_first_projections`) + # the cuBLAS first projection multiplies: the real `[down; block_inject]` rows + # rounded up to 32 (352 here). GB300, fp32 out, K 10240: 16384 rows 81.5 us + # against 91.5 at the padded 384, 8193 rows 41.6 against 51.2, never slower + # below; the real columns are bitwise the 384-wide product's at every row + # count measured (1-16384; at 320 columns they are not). + self.hc_first_cols = -(-(self.hc_lowrank + self.hc_count) // 32) * 32 # ---- full attention --------------------------------------------- self.heads_q = cfg.num_attention_heads @@ -1093,6 +1205,19 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig # Packed [q | k | v] scratch for the fused norm+rope call; the spare KV head # is the v slot the op needs but never reads. self.index_scratch_width = self.index_width + self.index_head_dim + # The fused input projection's output row (`weights._attn_in_w`): + # [q | k | v | gate | index q, k | v slot]. Its leading columns are the + # packed QKV the main norm+rope rewrites in place, and its tail is the + # index scratch's layout, so the row is `qkv_width + q_width + + # index_scratch_width` wide. The norm+rope addresses a row as + # `(Hq + Hk + Hv) * head_dim` and never touches the v heads, so the + # row must be a whole number of heads: everything after k is "v" to it. + self.attn_in_width = self.qkv_width + self.q_width + self.index_scratch_width + assert self.attn_in_width % self.head_dim == 0, ( + f"the fused projection row ({self.attn_in_width}) is not a whole number of " + f"{self.head_dim}-wide heads, so the main norm+rope cannot address it in place" + ) + self.attn_in_tail_heads = self.attn_in_width // self.head_dim - self.heads_q - self.heads_kv self.softmax_scale = 1.0 / math.sqrt(self.head_dim) # ---- gated deltanet ---------------------------------------------- @@ -1110,6 +1235,8 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig self.key_dim = self.num_k_heads * self.head_k_dim self.value_dim = self.num_v_heads * self.head_v_dim self.conv_dim = 2 * self.key_dim + self.value_dim + # Column widths of the one input-projection GEMM's output: [q|k|v | z | a | b]. + self.in_proj_split = (self.conv_dim, self.value_dim, self.num_v_heads, self.num_v_heads) # Gate activation of the gated RMS norm (`silu` or `sigmoid`). gate_type = getattr(cfg, "output_gate_type", None) assert gate_type == "sigmoid", ( @@ -1362,6 +1489,7 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig # Derived once after load (`post_load_weights`). self._moe_scalars: list | None = None + self._moe_scalars_cd: list | None = None self._gdn_gating: dict | None = None self._cos_sin: torch.Tensor | None = None # The attention op's KV scales: fp32 [1] ones over an fp8 pool (never None there: the @@ -1384,6 +1512,9 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig ) self._prefill: "PrefillTarget | None" = None self._decode: "DecodeTarget | None" = None + # The shared expert's stream; see `_moe_block`. Made in `post_load_weights`: + # a stream created inside a graph capture is not capturable. + self._side_stream: torch.cuda.Stream | None = None self.vision_tower: VisionTower | None = None # Everything derived from the KV cache manager (`_bind_cache_manager`); # None until the first forward. @@ -1400,6 +1531,41 @@ def _norm_slices(self, key: str) -> tuple: """Split a grouped-norm weight into its per-group slices (taken once).""" return tuple(torch.split(self.w[key], self.hidden, dim=0)) + def _fuse_hc_first_projections(self) -> None: + """Concatenate each trunk hyper-connection's `down` and `block_inject` + weights into one `[rows, stream]` tensor, zero-padded to a multiple of + 128 rows, so a site's first projection is one GEMM + (`_hc_first[(layer, tag)]`; `_hc_first["fhc"]` is the stack-final + mixer's `down`, padded the same way). cuBLAS was measured 1.8x (decode) + to 4.5x (8192 rows) faster at N = 384 than at the unpadded 324; its calls + multiply the first `hc_first_cols` (352) columns, faster still. + + The checkpoint's parameters are rebound as row views of the padded + tensors, so this costs only the padding rows (~1.2 MB per site), and + anything still reading them separately reads the same values.""" + w = self.w + + def fuse(*params): + rows = sum(t.shape[0] for t in params) + padded = -(-rows // 128) * 128 + pad = torch.zeros( + [padded - rows, self.stream], dtype=params[0].dtype, device=params[0].device + ) + both = torch.cat([t.data for t in params] + [pad], 0) + start = 0 + for t in params: + n = t.shape[0] + t.data = both[start : start + n] + start += n + return both + + self._hc_first = { + (i, tag): fuse(w[f"l{i}_{tag}_down"], w[f"l{i}_{tag}_inject"]) + for i in range(self.num_layers) + for tag in ("ahc", "mhc") + } + self._hc_first["fhc"] = fuse(w["fhc_down"]) + def post_load_weights(self) -> None: """Derive what the targets read, then build them, now that the weights are real. @@ -1409,7 +1575,15 @@ def post_load_weights(self) -> None: moe_scalars = [ _moe_layer_scalars(w, i, self.trunk_moe_backend) for i in range(self.num_layers) ] + # Under trtllm-gen, the CuteDSL scalars of the second expert copy that + # prefill-sized launches run on (`_CUTE_DSL_MOE_MIN_ROWS`). + moe_scalars_cd = ( + [_moe_layer_scalars(w, i, "CUTEDSL") for i in range(self.num_layers)] + if self.trunk_moe_backend == "TRTLLM" + else None + ) self._moe_scalars = moe_scalars + self._moe_scalars_cd = moe_scalars_cd # Both delta-rule entries index A_log and dt_bias as `ptr + head` and certify # fp32; the checkpoint is bf16. self._gdn_gating = { @@ -1455,6 +1629,10 @@ def post_load_weights(self) -> None: if self.mtp_enabled: self._mtp = _mtp_views(self) + self._fuse_hc_first_projections() + + # The CPU-only load checks build the targets over host weights; they launch nothing. + self._side_stream = torch.cuda.Stream(device=device) if device.type == "cuda" else None self._prefill = PrefillTarget(self) self._decode = DecodeTarget(self) self.vision_tower = VisionTower(self) @@ -1889,27 +2067,19 @@ class _LayerWeights(NamedTuple): """One trunk layer's weights as the views the ops read (GEMM operands are `.t()` views). Attention fields are None on a Gated-DeltaNet layer and vice versa.""" - ahc_norm: tuple - ahc_down: torch.Tensor + ahc_norm: torch.Tensor + ahc_first: torch.Tensor ahc_up: torch.Tensor - ahc_inject: torch.Tensor - mhc_norm: tuple - mhc_down: torch.Tensor + mhc_norm: torch.Tensor + mhc_first: torch.Tensor mhc_up: torch.Tensor - mhc_inject: torch.Tensor - q: torch.Tensor | None - k: torch.Tensor | None - v: torch.Tensor | None + attn_in: torch.Tensor | None o: torch.Tensor | None q_norm: torch.Tensor | None k_norm: torch.Tensor | None - idx: torch.Tensor | None idx_qn: torch.Tensor | None idx_kn: torch.Tensor | None - in_qkv: torch.Tensor | None - in_z: torch.Tensor | None - in_a: torch.Tensor | None - in_b: torch.Tensor | None + in_proj: torch.Tensor | None out: torch.Tensor | None conv_w: torch.Tensor | None a_log: torch.Tensor | None @@ -1929,9 +2099,21 @@ class _LayerWeights(NamedTuple): alpha1: torch.Tensor | None alpha2: torch.Tensor | None global_sf: torch.Tensor | None + # Under trtllm-gen, the second expert copy in the CuteDSL layout and its scalars, + # which prefill-sized launches run on (`_CUTE_DSL_MOE_MIN_ROWS`); None on CuteDSL. + fc1_w_cd: torch.Tensor | None + fc1_s_cd: torch.Tensor | None + fc2_w_cd: torch.Tensor | None + fc2_s_cd: torch.Tensor | None + alpha1_cd: torch.Tensor | None + alpha2_cd: torch.Tensor | None + global_sf_cd: torch.Tensor | None sh_gu: torch.Tensor sh_dn: torch.Tensor sh_gate: torch.Tensor + # `[gate | up | gate logit | pad]` (`weights._cat_rows`): `sh_gu` and `sh_gate` are its + # column blocks. + sh_gug: torch.Tensor def _layer_views(core: Qwen4ExpTp1ModelingV2) -> list[_LayerWeights]: @@ -1946,21 +2128,15 @@ def _layer_views(core: Qwen4ExpTp1ModelingV2) -> list[_LayerWeights]: layers = [] for i in range(core.num_layers): attn: dict = dict( - q=None, - k=None, - v=None, + attn_in=None, o=None, q_norm=None, k_norm=None, - idx=None, idx_qn=None, idx_kn=None, ) gdn: dict = dict( - in_qkv=None, - in_z=None, - in_a=None, - in_b=None, + in_proj=None, out=None, conv_w=None, a_log=None, @@ -1969,23 +2145,17 @@ def _layer_views(core: Qwen4ExpTp1ModelingV2) -> list[_LayerWeights]: ) if core.is_attn[i]: attn = dict( - q=w[f"l{i}_q_w"].t(), - k=w[f"l{i}_k_w"].t(), - v=w[f"l{i}_v_w"].t(), + attn_in=w[f"l{i}_attn_in_w"].t(), o=w[f"l{i}_o_w"].t(), q_norm=w[f"l{i}_q_norm"], k_norm=w[f"l{i}_k_norm"], - idx=w[f"l{i}_idx_w"].t(), idx_qn=w[f"l{i}_idx_qn"], idx_kn=w[f"l{i}_idx_kn"], ) else: a_log, dt_bias = core._gdn_gating[i] gdn = dict( - in_qkv=w[f"l{i}_in_qkv_w"].t(), - in_z=w[f"l{i}_in_z_w"].t(), - in_a=w[f"l{i}_in_a_w"].t(), - in_b=w[f"l{i}_in_b_w"].t(), + in_proj=w[f"l{i}_in_proj_w"].t(), out=w[f"l{i}_out_w"].t(), # [conv_dim, 1, width] -> [conv_dim, width]: a view (unit width stride). conv_w=w[f"l{i}_conv_w"].reshape(core.conv_dim, core.conv_kernel), @@ -1994,16 +2164,34 @@ def _layer_views(core: Qwen4ExpTp1ModelingV2) -> list[_LayerWeights]: gnorm=w[f"l{i}_gnorm"], ) scalars = core._moe_scalars[i] + cute_dsl_copy: dict = dict( + fc1_w_cd=None, + fc1_s_cd=None, + fc2_w_cd=None, + fc2_s_cd=None, + alpha1_cd=None, + alpha2_cd=None, + global_sf_cd=None, + ) + if core._moe_scalars_cd is not None: + scalars_cd = core._moe_scalars_cd[i] + cute_dsl_copy = dict( + fc1_w_cd=w[f"l{i}_fc1_w_cd"], + fc1_s_cd=w[f"l{i}_fc1_s_cd"], + fc2_w_cd=w[f"l{i}_fc2_w_cd"], + fc2_s_cd=w[f"l{i}_fc2_s_cd"], + alpha1_cd=scalars_cd["alpha1"], + alpha2_cd=scalars_cd["alpha2"], + global_sf_cd=scalars_cd["global_sf"], + ) layers.append( _LayerWeights( - ahc_norm=core._norm_slices(f"l{i}_ahc_norm"), - ahc_down=w[f"l{i}_ahc_down"], - ahc_up=w[f"l{i}_ahc_up"], - ahc_inject=w[f"l{i}_ahc_inject"], - mhc_norm=core._norm_slices(f"l{i}_mhc_norm"), - mhc_down=w[f"l{i}_mhc_down"], - mhc_up=w[f"l{i}_mhc_up"], - mhc_inject=w[f"l{i}_mhc_inject"], + ahc_norm=w[f"l{i}_ahc_norm"], + ahc_first=core._hc_first[(i, "ahc")].t(), + ahc_up=w[f"l{i}_ahc_up"].t(), + mhc_norm=w[f"l{i}_mhc_norm"], + mhc_first=core._hc_first[(i, "mhc")].t(), + mhc_up=w[f"l{i}_mhc_up"].t(), **attn, **gdn, router=w[f"l{i}_router"].t(), @@ -2018,9 +2206,11 @@ def _layer_views(core: Qwen4ExpTp1ModelingV2) -> list[_LayerWeights]: alpha1=scalars.get("alpha1"), alpha2=scalars.get("alpha2"), global_sf=scalars.get("global_sf"), - sh_gu=w[f"l{i}_sh_gu"].t(), + **cute_dsl_copy, + sh_gu=w[f"l{i}_sh_gu"][: 2 * core.shared_inter].t(), sh_dn=w[f"l{i}_sh_dn"].t(), - sh_gate=w[f"l{i}_sh_gate"].t(), + sh_gate=w[f"l{i}_sh_gu"][2 * core.shared_inter : 2 * core.shared_inter + 1].t(), + sh_gug=w[f"l{i}_sh_gu"].t(), ) ) return layers @@ -2030,14 +2220,17 @@ class _StackWeights(NamedTuple): """The weights outside the layer loop: the stack-final mixer (this checkpoint's final norm) and the PLE side path.""" - fhc_norm: tuple - fhc_down: torch.Tensor + fhc_norm: torch.Tensor + fhc_first: torch.Tensor fhc_up: torch.Tensor ple_key: torch.Tensor ple_val: torch.Tensor ple_nk: tuple ple_nq: tuple ple_nc: tuple + ple_nk_w: torch.Tensor + ple_nq_w: torch.Tensor + ple_nc_w: torch.Tensor ple_conv_w: torch.Tensor ple_mult: torch.Tensor ple_vocab: torch.Tensor @@ -2048,14 +2241,19 @@ class _StackWeights(NamedTuple): def _stack_views(core: Qwen4ExpTp1ModelingV2) -> _StackWeights: w = core.w return _StackWeights( - fhc_norm=core._norm_slices("fhc_norm"), - fhc_down=w["fhc_down"], - fhc_up=w["fhc_up"], + fhc_norm=w["fhc_norm"], + fhc_first=core._hc_first["fhc"].t(), + fhc_up=w["fhc_up"].t(), ple_key=w["ple_key_w"].t(), ple_val=w["ple_val_w"].t(), ple_nk=core._norm_slices("ple_nk"), ple_nq=core._norm_slices("ple_nq"), ple_nc=core._norm_slices("ple_nc"), + # The same three weights whole, for the decode target's fused tail, which + # normalizes each stream inside one program. + ple_nk_w=w["ple_nk"], + ple_nq_w=w["ple_nq"], + ple_nc_w=w["ple_nc"], ple_conv_w=w["ple_conv_w"], ple_mult=w["ple_mult"], ple_vocab=w["ple_vocab"], @@ -2132,9 +2330,9 @@ def _mtp_views(core: Qwen4ExpTp1ModelingV2) -> _MtpWeights: fc1_s=w[f"{p}fc1_s"], fc2_w=w[f"{p}fc2_w"], fc2_s=w[f"{p}fc2_s"], - sh_gu=w[f"{p}sh_gu"].t(), + sh_gu=w[f"{p}sh_gu"][: 2 * core.shared_inter].t(), sh_dn=w[f"{p}sh_dn"].t(), - sh_gate=w[f"{p}sh_gate"].t(), + sh_gate=w[f"{p}sh_gu"][2 * core.shared_inter : 2 * core.shared_inter + 1].t(), fhc_norm=core._norm_slices(f"{p}fhc_norm"), fhc_down=w[f"{p}fhc_down"], fhc_up=w[f"{p}fhc_up"], @@ -2236,6 +2434,80 @@ def _sparse_gqa( ) +def _attn_in_heads(core, proj): + """The main q, k and v heads of the fused attention projection's output `proj` + `[T, attn_in_width]` (`weights._attn_in_w`), as `[T, heads, head_dim]` views.""" + q, k, v, _ = torch.split( + torch.reshape(proj, [proj.shape[0], core.attn_in_width // core.head_dim, core.head_dim]), + [core.heads_q, core.heads_kv, core.heads_kv, core.attn_in_tail_heads - core.heads_kv], + 1, + ) + return q, k, v + + +def _cute_dsl_routed(core, L, logits, xq, xsf, routed_out) -> torch.Tensor: + """The trunk's routed experts on the CuteDSL pipeline: Renormalize routing, the sort + into per-expert row tiles, FC1 gathering its rows from the unpermuted NVFP4 activation + `xq` / `xsf` (linear scales) with SwiGLU and the requantize to the layer's one FC2 input + scale fused, FC2, and the routing-weighted unpermute in a fixed slot order, which + writes every row of `routed_out` `[T, hidden]` and returns it.""" + # fp32 weights: the dtype the unpermute scales by. + topk_ids, topk_w = renorm_moe_routing_op(logits, topk=core.topk, output_dtype=torch.float32) + # The sort pads each expert's (token, expert) pairs to whole tiles and emits the tile + # tables both GEMMs walk. + tile_expert, tile_limit, expanded_to_permuted, permuted_to_expanded, _, live_tiles = moe_sort( + token_selected_experts=topk_ids, + token_final_scales=topk_w, + num_experts=core.num_experts, + top_k=core.topk, + local_expert_offset=0, + local_num_experts=core.num_experts, + tile_tokens_dim=_CUTEDSL_TILE_SIZE, + ) + h, h_sf = cute_dsl_nvfp4_gather_grouped_gemm_act_fusion( + input=xq.view(torch.float4_e2m1fn_x2), + weight=L.fc1_w, + input_scale=xsf.view(torch.uint8).view(xq.shape[0], core.hidden // _SF_VEC), + weight_scale=L.fc1_s, + alpha=L.alpha1, + tile_idx_to_group_idx=tile_expert, + tile_idx_to_mn_limit=tile_limit, + permuted_idx_to_expanded_idx=permuted_to_expanded, + num_non_exiting_tiles=live_tiles, + global_sf=L.global_sf, + num_experts=core.num_experts, + top_k=core.topk, + num_local_experts=core.num_experts, + local_expert_offset=0, + tile_size=_CUTEDSL_TILE_SIZE, + scaling_vector_size=_SF_VEC, + activation_type=int(ActivationType.Swiglu), + ) + # FC2 writes one bf16 row per permuted row. + y = cute_dsl_nvfp4_grouped_gemm( + input=h, + weight=L.fc2_w, + input_scale=h_sf, + weight_scale=L.fc2_s, + alpha=L.alpha2, + tile_idx_to_group_idx=tile_expert, + num_non_exiting_tiles=live_tiles, + num_experts=core.num_experts, + top_k=core.topk, + num_local_experts=core.num_experts, + local_expert_offset=0, + tile_size=_CUTEDSL_TILE_SIZE, + output_dtype=torch.bfloat16, + scaling_vector_size=_SF_VEC, + ) + return moe_unpermute_inplace( + permuted_input=y, + output=routed_out, + expanded_idx_to_permuted_idx=expanded_to_permuted, + topk_scales=topk_w, + ) + + def _live_score_columns(core, md, batch: int) -> int: """A host-side bound on every row's selection window in a step with context rows: the longest sequence's compressed length, with the verify width as slack for the draft @@ -2446,6 +2718,38 @@ def step_args(self, md: TrtllmAttentionMetadata) -> dict: # --------------------------------------------------------------- blocks --- + def _hc_mixer(self, normed, first, up, with_inject): + """One hyper-connection mixer from the site's normed stream: the first + projection over `[down; block_inject; 0]` (fp32 logits), the SiLU and + injection activations, then the up projection with the sigmoid gate and + the stream mean in its epilogue (the gate logits never leave the GEMM). + Returns `(mixed, inject)`; `inject` is None without `with_inject` (the + stack-final mixer).""" + core = self.core + low, inject = qwen4_exp_hc_lowrank_act( + cublas_mm(normed, first[:, : core.hc_first_cols], None, torch.float32), + core.hc_lowrank, + core.hc_count, + with_inject, + normed.dtype, + ) + return qwen4_exp_hc_up_gate_mean(low, up, normed, core.hc_count), inject + + def _hc_moe_mixer(self, normed, first, up, gscale): + """`_hc_mixer` at an MoE site (always with injection) whose fused up + projection also writes the routed experts' NVFP4 input (linear scales, + the layer's global scale `gscale`). Returns `(mixed, inject, xq, xsf)`.""" + core = self.core + low, inject = qwen4_exp_hc_lowrank_act( + cublas_mm(normed, first[:, : core.hc_first_cols], None, torch.float32), + core.hc_lowrank, + core.hc_count, + True, + normed.dtype, + ) + mixed, xq, xsf = qwen4_exp_hc_up_gate_mean_fp4(low, up, normed, core.hc_count, gscale) + return mixed, inject, xq, xsf + def _grouped_norm(self, x: torch.Tensor, weights: tuple) -> torch.Tensor: """`Qwen4ExpTextRMSNorm(group_size=hidden)`: an independent RMS per `hidden`-wide group, then `(1 + w)` over the full width; one gemma-RMS-norm per column @@ -2456,209 +2760,144 @@ def _grouped_norm(self, x: torch.Tensor, weights: tuple) -> torch.Tensor: [flashinfer_gemma_rmsnorm(part, w, core.eps) for w, part in zip(weights, parts)], 1 ) - def _moe_block( - self, x: torch.Tensor, L: _LayerWeights, dt: torch.dtype, routed_out: torch.Tensor | None - ) -> 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. `routed_out` - is the CuteDSL path's `[T, hidden]` output buffer (None on trtllm-gen, whose - runner allocates its own).""" + def _moe_block(self, x, xq, xsf, L: _LayerWeights, routed_out) -> tuple: + """One MoE layer: NVFP4 routed experts plus a bf16 shared expert, left + uncombined. Returns `(routed, (gate_logits, shared))`: the next + `qwen4_exp_hc_inject_norm` forms `routed + sigmoid(gate_logits) * shared` + inside its kernel. + + The shared expert and its per-token gate logit run on the core's side + stream, concurrently with the routed path, and join it before this + returns, so every consumer on the main stream is ordered after both. + `x` is the bf16 hidden states (router and shared expert); `xq` / `xsf` + are the same states in NVFP4 with **linear** block scales, written by the + MoE-site mixer (`_hc_moe_mixer`). On trtllm-gen the routed path hands the + runner the router logits (its Renormalize routing is bitwise the + `renorm_moe_routing_op` + pre-routed composition), and a launch of at + least `_CUTE_DSL_MOE_MIN_ROWS` rows runs the CuteDSL MoE over the second + expert copy instead; on CuteDSL every launch is `_cute_dsl_routed`, which + writes `routed_out`.""" core = self.core - # The routing op reads logits from data_ptr() ignoring strides; cublas_mm - # allocates them contiguous. + main = torch.cuda.current_stream() + side = core._side_stream + # Fork: the side stream reads `x` only after everything queued on `main` + # so far. What it allocates is consumed on `main` after the join below, + # and the next layer's fork orders any reuse of those blocks after that. + side.wait_stream(main) + with torch.cuda.stream(side): + # One GEMM forms gate|up and the gate logit (`sh_gug`'s extra rows): a separate + # gate GEMV at prefill sizes keeps the SMs from the routed path's persistent + # FC1 until it drains. + gug = cublas_mm(x, L.sh_gug) + shared = cublas_mm(silu_and_mul_rows(gug, core.shared_inter), L.sh_dn) + gate_logits = gug[:, 2 * core.shared_inter : 2 * core.shared_inter + 1] + # The routing reads its logits as a dense row-major buffer from data_ptr() + # and ignores strides; cublas_mm allocates them contiguous. logits = cublas_mm(x, L.router) if core.trunk_moe_backend == "CUTEDSL": - # fp32 weights: the dtype the unpermute scales by. - topk_ids, topk_w = renorm_moe_routing_op( - logits, topk=core.topk, output_dtype=torch.float32 - ) - xq, xsf = fp4_quantize(x, L.moe_gscale, _SF_VEC, False, _SF_LINEAR) - # The sort pads each expert's (token, expert) pairs to whole tiles and emits - # the tile tables both GEMMs walk. - tile_expert, tile_limit, expanded_to_permuted, permuted_to_expanded, _, live_tiles = ( - moe_sort( - token_selected_experts=topk_ids, - token_final_scales=topk_w, - num_experts=core.num_experts, - top_k=core.topk, - local_expert_offset=0, - local_num_experts=core.num_experts, - tile_tokens_dim=_CUTEDSL_TILE_SIZE, - ) - ) - # FC1 gathers its rows from the unpermuted activation, applies SwiGLU and - # requantizes with the layer's one FC2 input scale. `fp4_quantize` returns the - # linear scales flat; the GEMM reads them as [T, hidden / 16]. - h, h_sf = cute_dsl_nvfp4_gather_grouped_gemm_act_fusion( - input=xq.view(torch.float4_e2m1fn_x2), - weight=L.fc1_w, - input_scale=xsf.view(x.shape[0], core.hidden // _SF_VEC), - weight_scale=L.fc1_s, - alpha=L.alpha1, - tile_idx_to_group_idx=tile_expert, - tile_idx_to_mn_limit=tile_limit, - permuted_idx_to_expanded_idx=permuted_to_expanded, - num_non_exiting_tiles=live_tiles, - global_sf=L.global_sf, - num_experts=core.num_experts, - top_k=core.topk, - num_local_experts=core.num_experts, - local_expert_offset=0, - tile_size=_CUTEDSL_TILE_SIZE, - scaling_vector_size=_SF_VEC, - activation_type=int(ActivationType.Swiglu), - ) - # FC2 writes one bf16 row per permuted row; the unpermute writes every row of - # `routed_out` as the routing-weighted sum of its token's top-k rows. - y = cute_dsl_nvfp4_grouped_gemm( - input=h, - weight=L.fc2_w, - input_scale=h_sf, - weight_scale=L.fc2_s, - alpha=L.alpha2, - tile_idx_to_group_idx=tile_expert, - num_non_exiting_tiles=live_tiles, - num_experts=core.num_experts, - top_k=core.topk, - num_local_experts=core.num_experts, - local_expert_offset=0, - tile_size=_CUTEDSL_TILE_SIZE, - output_dtype=torch.bfloat16, - scaling_vector_size=_SF_VEC, - ) - routed = moe_unpermute_inplace( - permuted_input=y, - output=routed_out, - expanded_idx_to_permuted_idx=expanded_to_permuted, - topk_scales=topk_w, + routed = _cute_dsl_routed(core, L, logits, xq, xsf, routed_out) + elif x.shape[0] >= _CUTE_DSL_MOE_MIN_ROWS: + routed = fp4_moe_renorm_cute_dsl( + logits, + xq, + xsf, + L.fc1_w_cd, + L.fc1_s_cd, + L.alpha1_cd, + L.fc2_w_cd, + L.fc2_s_cd, + L.alpha2_cd, + L.global_sf_cd, + core.num_experts, + core.topk, ) else: - 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) + routed = fp4_block_scale_moe_renorm_routed( + logits, + xq, + xsf.view(torch.float8_e4m3fn), + L.fc1_w, + L.fc1_s, + L.fc2_w, + L.fc2_s, + L.out1_scale, + L.out1_gate_scale, + L.out2_scale, + core.num_experts, + core.topk, + core.moe_inter, + ) + main.wait_stream(side) + return routed, (gate_logits, shared) - def _attention_block(self, x, i, L, qkv, index_scratch, attn_out, pos, pos3): + def _attention_block(self, x, i, L, proj_qkv, index_scratch, attn_out, pos3): """One gated GQA layer plus its QSA index branch. - The q projection is twice as wide as q (q and output gate interleave per - head, as HF chunks the last dim); only the q half is packed into the QKV - buffer and the gate half is read later through its own strides. The - projections, the packing and both ropes are row-wise and run over all `T` rows - here; the cache writes, the selection and the attention core read the step's - batch and run in the `_attention_core` break, which leaves the attention output - in `attn_out`.""" + The four input projections are one GEMM against the load-time fused + operand (`weights._attn_in_w`), whose output row is + `[q | k | v | gate | index q, k | v slot]`. One norm+rope launch rewrites + the main q/k heads in place on that output and writes the normalized, + rotated index query into the scratch; v, the gate and the raw index key + are not touched. Both are row-wise and run over all `T` rows here; the + cache writes, the selection and the attention core read the step's batch + and run in the `_attention_core` break, which leaves the attention output + in `attn_out`. The gate is read through its row stride.""" core = self.core num_tokens = x.shape[0] - 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, + proj = cublas_mm(x, L.attn_in) + gate = torch.split(proj, [core.qkv_width, core.q_width, core.index_scratch_width], 1)[1] + # The index query lands in the q slots of the scratch; its k and v slots + # are layout only and nothing reads them. + q_index = torch.split( + torch.reshape( + index_scratch, + [num_tokens, core.index_heads + 2 * core.index_kv_heads, core.index_head_dim], + ), + [core.index_heads, core.index_kv_heads, core.index_kv_heads], + 1, + )[0] + # Both planes' Gemma norm + rope in one launch, over the target's fp32 + # cos|sin table: the main q/k heads in place (the real interleaved mRoPE; + # three agreeing rows reduce exactly to the plain rope on row 0), the index + # query from the projection's index columns into `q_index` (coordinate 0 + # alone, matching `_QSA_MROPE_SECTION` on the key side). + index_section = (_QSA_MROPE_SECTION or (0, 0, 0)) if _QSA_INDEX_USE_MROPE else (0, 0, 0) + qwen4_exp_qk_norm_rope( + proj, + q_index, + pos3, + core._cos_sin, + L.q_norm, + L.k_norm, + L.idx_qn, num_heads_q=core.heads_q, num_heads_k=core.heads_kv, - num_heads_v=core.heads_kv, head_dim=core.head_dim, + index_heads=core.index_heads, + index_head_dim=core.index_head_dim, + index_q_col=core.qkv_width + core.q_width, rotary_dim=core.rotary_dim, eps=core.eps, - q_weight=L.q_norm, - k_weight=L.k_norm, - base=core.theta, - is_neox=True, - position_ids=pos3, - is_qk_norm=True, - # The gemma QK norm, `rmsnorm(x) * (1 + w)`. - use_gemma=True, - # Real mRoPE, unconditionally. 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, - ) - if self._kv_probe is not None and not self._kv_probe.observe(i, k_dst, v_dst): - self._kv_probe = None - - # One replicated projection produces [index_heads q | 1 k]; its raw key is what - # the break stores. The query's norm and rotation run here in the fused call's - # packed layout, on a scratch copy: 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. - qk_index = cublas_mm(x, L.idx) - 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], - ) - eager_on_graph(qkv, qk_index, index_scratch, attn_out, i, inner=self._attention_core) + index_mrope_section1=index_section[1], + index_mrope_section2=index_section[2], + ) + if self._kv_probe is not None: + _, k_dst, v_dst = _attn_in_heads(core, proj) + if not self._kv_probe.observe(i, k_dst, v_dst): + self._kv_probe = None + eager_on_graph(proj, proj_qkv, index_scratch, attn_out, i, inner=self._attention_core) fused_sigmoid_mul_inplace(attn_out, gate) return cublas_mm(attn_out, L.o) - def _attention_core(self, qkv, qk_index, index_scratch, attn_out, i) -> None: + def _attention_core(self, proj, proj_qkv, index_scratch, attn_out, i) -> None: """Break: the index plane's store and compression, the dense/sparse decision and the attention core over the step's `n` rows. Writes rows `[:n]` of `attn_out` - `[T, q_width]` in place and zeroes rows `[n:]`. + `[T, q_width]` in place and zeroes rows `[n:]`. `proj` is the layer's fused + projection output after the norm+rope, `proj_qkv` the dense op's packed QKV + scratch. The index plane advances on every step, including below the threshold, so a request growing past it finds its prefix compressed; above it the main K/V append @@ -2686,16 +2925,19 @@ def _attention_core(self, qkv, qk_index, index_scratch, attn_out, i) -> None: live_cols, ) = step.qsa num_tokens = step.num_tokens - qkv = qkv[:num_tokens] - 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) + proj = proj[:num_tokens] + qkv_view, _, index_view = torch.split( + proj, [core.qkv_width, core.q_width, core.index_scratch_width], 1 + ) + q_dst, k_dst, v_dst = _attn_in_heads(core, proj) index_cache = core._bound.index_cache[i] token_k = torch.split( torch.reshape( - qk_index[:num_tokens], [num_tokens, core.index_heads + core.index_kv_heads, -1] + index_view, + [num_tokens, core.index_scratch_width // core.index_head_dim, core.index_head_dim], ), - [core.index_heads, core.index_kv_heads], + [core.index_heads, core.index_kv_heads, 1], 1, )[1] if step.stage: @@ -2740,44 +2982,100 @@ def _attention_core(self, qkv, qk_index, index_scratch, attn_out, i) -> None: # unit KV scales; over an NVFP4 pool they quantize and dequantize at the dense op's # second-level scales. _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block) - 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, + # The step's context rows come first, consecutive rows of a request, and score + # as query tiles; its generation rows, one a request, score row by row (a query + # tile of them would walk every request's keys once per row). Each part selects + # into its own rows of the top-k. + num_ctx = step.attn["num_ctx_tokens"] + parts = [(0, num_ctx, True)] + if num_ctx < num_tokens: + parts.append((num_ctx, num_tokens, False)) + # Each part's canonical block top-k and its expansion are one launch writing its + # rows of the step's selection (QSA's row starts are zero); a score table wider + # than that kernel's 4096 columns takes the composition. + selected = torch.empty( + [num_tokens, core.token_topk + core.compress_ratio - 1], + dtype=torch.int32, + device=proj.device, ) - _select(core, scores, topk_row_starts, visible_blocks, topk_indices, live_cols) - selected = qsa_expand_block_indices( - topk_indices, - logical, - seq_lens_q, - compress_ratio=core.compress_ratio, - token_topk=core.token_topk, - ) - attn_out[:num_tokens].copy_( - torch.reshape( - _sparse_gqa( + fused_select = block_table.shape[1] * tokens_per_block // core.compress_ratio <= 4096 + for lo, hi, context_rows in parts: + scores = qsa_paged_index_scores( + q=q_index[lo:hi], + index_cache=index_cache, + block_table=block_table, + query_positions=logical[lo:hi], + request_indices=req_idx[lo:hi], + tokens_per_block=tokens_per_block, + compress_ratio=core.compress_ratio, + only_visible_blocks=True, + context_rows=context_rows, + ) + if fused_select and scores.shape[1] % 4 == 0: + qsa_topk_expand_decode( + scores, + visible_blocks[lo:hi], + logical[lo:hi], + seq_lens_q[lo:hi], + topk_indices[lo:hi], + core.compress_ratio, + core.token_topk, + out=selected[lo:hi], + ) + else: + _select( core, - kv, - q=q_dst, - block_table=block_table, - selected_tokens=selected, - request_indices=req_idx, - tokens_per_block=tokens_per_block, - # A bound, not a mask: lets the kernel stop scanning a row early - # (context-batch branch only). - query_positions=logical, - compress_ratio=core.compress_ratio, - ), - [num_tokens, core.q_width], + scores, + topk_row_starts[lo:hi], + visible_blocks[lo:hi], + topk_indices[lo:hi], + live_cols, + ) + selected[lo:hi].copy_( + qsa_expand_block_indices( + topk_indices[lo:hi], + logical[lo:hi], + seq_lens_q[lo:hi], + compress_ratio=core.compress_ratio, + token_topk=core.token_topk, + ) + ) + # A context launch of thousands of rows over a bf16 pool takes the row-tile + # pass: consecutive rows of one request share their K/V tiles; rows of groups + # that span requests, the generation rows of a mixed step among them, are served + # row by row. Below the threshold the per-row kernel is faster. + if num_tokens >= _ROWTILE_MIN_ROWS and kv.k.dtype == q_dst.dtype: + out = qsa_sparse_gqa_rowtile( + q=q_dst, + k_cache=kv.k, + v_cache=kv.v, + block_table=block_table, + selected_tokens=selected, + request_indices=req_idx, + tokens_per_block=tokens_per_block, + softmax_scale=core.softmax_scale, + query_positions=logical, + compress_ratio=core.compress_ratio, ) - ) + else: + out = _sparse_gqa( + core, + kv, + q=q_dst, + block_table=block_table, + selected_tokens=selected, + request_indices=req_idx, + tokens_per_block=tokens_per_block, + # A bound, not a mask: lets the kernel stop scanning a row early + # (context-batch branch only). + query_positions=logical, + compress_ratio=core.compress_ratio, + ) + attn_out[:num_tokens].copy_(torch.reshape(out, [num_tokens, core.q_width])) else: + # The dense op reads its packed QKV flat: one compact copy. + qkv = proj_qkv[:num_tokens] + qkv.copy_(qkv_view) thop_attention( q=qkv, output=attn_out[:num_tokens], @@ -2883,23 +3181,25 @@ def _gdn_verify( return torch.reshape(out, [1, n * p, core.num_v_heads, core.head_v_dim]) def _gdn_block(self, x, i, L, core_out) -> torch.Tensor: - """One Gated-DeltaNet layer over a batch with a context prefix: the four input - projections and the gated-norm epilogue run over all `T` rows here, the - recurrence in the `_gdn_core` break, which leaves its output in `core_out` - `[1, T, num_v_heads, head_v_dim]`.""" - 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) - eager_on_graph(mixed_qkv, a, b, core_out, i, inner=self._gdn_core) + """One Gated-DeltaNet layer over a batch with a context prefix: the one input + projection GEMM (q|k|v, z, a and b are column views of its output) and the + gated-norm epilogue run over all `T` rows here, the recurrence in the `_gdn_core` + break, which leaves its output in `core_out` `[1, T, num_v_heads, head_v_dim]`.""" + core = self.core + proj = cublas_mm(x, L.in_proj) + z = torch.split(proj, core.in_proj_split, 1)[1] + eager_on_graph(proj, core_out, i, inner=self._gdn_core) return self._gdn_epilogue(core_out, z, L, x.shape[0]) - def _gdn_core(self, mixed_qkv, a, b, out, i) -> None: - """Break: the recurrence over the step's `n` rows (the conv, the post-conv gating and - the chunked delta rule, or under speculation the verify path for the generation - rows). Writes rows `[:n]` of `out` `[1, T, num_v_heads, head_v_dim]` in place and - zeroes rows `[n:]`. Every branch is on a host int of the step, never a tensor - value.""" + def _gdn_core(self, proj, out, i) -> None: + """Break: the recurrence over the step's `n` rows. Writes rows `[:n]` of `out` + `[1, T, num_v_heads, head_v_dim]` in place and zeroes rows `[n:]`. Every branch + is on a host int of the step, never a tensor value. + + Over an fp32 recurrent pool outside a verify step the context rows take one fused + conv + post-conv pass (`_gdn_prefill_prep`) and FlashInfer's chunked recurrence, and + a mixed step's one-token generation rows FlashInfer's single-token decode kernel; + otherwise `_gdn_core_generic`.""" core = self.core L = self._layers[i] step = self._step @@ -2922,9 +3222,119 @@ def _gdn_core(self, mixed_qkv, a, b, out, i) -> None: replay, ) = step.gdn num_tokens = step.num_tokens - mixed_qkv = mixed_qkv[:num_tokens] - a = a[:num_tokens] - b = b[:num_tokens] + ssm_states = core._bound.ssm[i] + if is_verify or ssm_states.dtype != torch.float32: + self._gdn_core_generic(proj, out, i) + return + mixed_qkv, _, a, b = torch.split(proj[:num_tokens], core.in_proj_split, 1) + conv_states = core._bound.conv[i] + + # A context sequence carrying nothing starts from a zeroed slot; only the + # context prefix is a candidate (a generation row's state must survive). + reset_recurrent_state_rows(ssm_states, si_p, hi_p, conv_states) + if num_decode_tokens > 0: + mq_p, mq_d = torch.split(mixed_qkv, [num_prefill_tokens, num_decode_tokens], 0) + a_p, b_p = a[:num_prefill_tokens], b[:num_prefill_tokens] + q, k, v, g, beta = self._gdn_prefill_prep(mq_p, a_p, b_p, L, i, qsl_p, si_p, hi_p) + causal_conv1d_update( + x=mq_d, + conv_state=conv_states, + weight=L.conv_w, + bias=None, + silu_activation=True, + cache_seqlens=None, + conv_state_indices=si_d, + pad_slot_id=_PAD_SLOT_ID, + ) + # The two halves run separate recurrences over disjoint slots and write + # disjoint row ranges of the output. + out_p = out[:, :num_prefill_tokens] + flashinfer_chunk_gated_delta_rule( + q, + k, + v, + g, + beta, + scale=None, + initial_state=ssm_states, + initial_state_indices=si_p, + inplace_indexed_state_update=True, + output_final_state=False, + cu_seqlens=qsl_long_p, + head_first=False, + use_qk_l2norm_in_kernel=False, + output=out_p, + state_workspace=None, + ) + # The generation rows: the decode step's recurrence (its kernel without the + # gated norm, which runs with the context rows' below) over their rows of the + # projection, whose q/k/v columns the conv step rewrote; it computes both gates + # and the q/k norm itself and reads the rows' pool slots before its PDL wait + # (the context rows' reset and recurrence touch only theirs). + gdn_decode_rows( + proj[num_prefill_tokens:num_tokens], + ssm_states, + si_d, + L.a_log, + L.dt_bias, + core.num_k_heads, + core.head_k_dim, + out[0, num_prefill_tokens:num_tokens].view(num_decode_tokens, core.value_dim), + ) + else: + q, k, v, g, beta = self._gdn_prefill_prep( + mixed_qkv, a, b, L, i, qsl_all, si_all, hi_all + ) + # q and k arrive L2-normalized; the recurrence must not repeat it. + flashinfer_chunk_gated_delta_rule( + q, + k, + v, + g, + beta, + scale=None, + initial_state=ssm_states, + initial_state_indices=si_all, + inplace_indexed_state_update=True, + output_final_state=False, + cu_seqlens=qsl_long, + head_first=False, + use_qk_l2norm_in_kernel=False, + output=out[:, :num_tokens], + state_workspace=None, + ) + out[:, num_tokens:].zero_() + + def _gdn_core_generic(self, proj, out, i) -> None: + """`_gdn_core` on a verify step or over a bf16 recurrent pool: the conv, the + post-conv gating and the chunked delta rule over the step's `n` rows, or under + speculation the verify path for the generation rows. The q|k|v, a and b columns are + copied out of the projection first: these entries are certified on contiguous + operands.""" + core = self.core + L = self._layers[i] + step = self._step + ( + num_prefill_tokens, + num_decode_tokens, + si_p, + si_all, + si_d, + hi_p, + hi_all, + qsl_p, + qsl_all, + qsl_long, + is_verify, + num_decodes, + gen_p, + int_idx, + qsl_long_p, + replay, + ) = step.gdn + num_tokens = step.num_tokens + mixed_qkv, _, a, b = torch.split(proj[:num_tokens], core.in_proj_split, 1) + mixed_qkv, a, b = mixed_qkv.contiguous(), a.contiguous(), b.contiguous() conv_states = core._bound.conv[i] ssm_states = core._bound.ssm[i] @@ -3062,6 +3472,32 @@ def _gdn_core(self, mixed_qkv, a, b, out, i) -> None: out[:, :num_tokens].copy_(recurrence) out[:, num_tokens:].zero_() + def _gdn_prefill_prep(self, mq, a, b, L, i, qsl, si, hi): + """The context rows' recurrence inputs: the causal conv against each sequence's + carried window (published back to the pool), SiLU, the q|k|v split with q/k + L2-normalized, and the decay/beta gates -- one pass over `mq`, the token-major + q|k|v column view of the projection.""" + core = self.core + return gdn_prefill_conv_post( + x=mq, + weight=L.conv_w, + conv_states=core._bound.conv[i], + query_start_loc=qsl, + cache_indices=si, + has_initial_state=hi, + a=a, + b=b, + A_log=L.a_log, + dt_bias=L.dt_bias, + num_k_heads=core.num_k_heads, + head_k_dim=core.head_k_dim, + num_v_heads=core.num_v_heads, + head_v_dim=core.head_v_dim, + pad_slot_id=_PAD_SLOT_ID, + l2_norm_eps=_L2_NORM_EPS, + softplus_threshold=_SOFTPLUS_THRESHOLD, + ) + def _gdn_epilogue(self, out, z, L, num_tokens) -> torch.Tensor: """Gated RMS norm and output projection shared by every recurrence.""" core = self.core @@ -3076,10 +3512,11 @@ def _gdn_epilogue(self, out, z, L, num_tokens) -> torch.Tensor: ) return cublas_mm(torch.reshape(o, [num_tokens, core.value_dim]), L.out) - def _ple_block(self, hc, out) -> None: - """Break: the n-gram side path on one decoder layer, over the step's `n` rows. Writes - rows `[:n]` of `out` `[T, stream]` in place and zeroes rows `[n:]`; the caller adds - `out` into the hyper-connection stream. + def _ple_block(self, hc, hc_out, normed_out) -> None: + """Break: the n-gram side path on one decoder layer, over the step's `n` rows, + added into the stream and normalized for the layer's attention-side site. Writes + rows `[:n]` of `hc_out` (the stream) and `normed_out` (its grouped norm), both + `[T, stream]`, in place and zeroes rows `[n:]`. Per-request stateful: ids become one `ngram_size`-wide window per token (predecessors from the paged id history); each window hashes to one row per @@ -3087,25 +3524,79 @@ def _ple_block(self, hc, out) -> None: scale; key/value projections and the gate scale one value into each stream; a dilated depthwise short conv over a paged state is added back. A sequence carrying nothing is seeded with `eos_token_id`, not 0: id 0 is a real token - the hash cannot distinguish. - - The id history and the conv state publish the state after the whole chunk. On - a verify step (`stage`) the generation rows' pre-step pool rows, id windows and - conv input go to `SpeculativeRollback`, which rewrites both pools to the - accepted prefix once the sampler has decided it.""" + the hash cannot distinguish. The id history and the conv state publish the + state after the whole chunk. + + Outside a verify step the window, hash, gather and dequantize are + `ple_ngram_embed`, and the gate, the conv, the residual add and the next site's + grouped norm `ple_gate_conv_inject_norm`. On a verify step (`stage`) the + generation rows' pre-step pool rows, id windows and conv input go to + `SpeculativeRollback`, which rewrites both pools to the accepted prefix once the + sampler has decided it, so that step takes `_ple_delta`, which exposes them.""" core = self.core S = self._stack step = self._step conv_states, ngram_pool, si_all, carry_all, qsl = step.ple - stage = step.stage num_tokens = step.num_tokens - input_ids = step.input_ids + next_norm = self._layers[core.ple_layer].ahc_norm hc = hc[:num_tokens] - - if stage: - core.spec_rollback.stage_ple_pools(ngram_pool, conv_states) + if step.stage: + stream = torch.add(hc, self._ple_delta(hc)) + stream, normed = qwen4_exp_hc_inject_norm( + stream, None, None, next_norm, core.eps, core.hc_count + ) + else: + ngram = ple_ngram_embed( + step.input_ids, + ngram_pool, + si_all, + carry_all, + qsl, + core.w["ple_table"], + S.ple_tscale, + S.ple_mult, + S.ple_vocab, + S.ple_offs, + core.eos_token_id, + num_tokens, + ) + stream, normed = ple_gate_conv_inject_norm( + hc, + cublas_mm(ngram, S.ple_key), + cublas_mm(ngram, S.ple_val), + S.ple_nk_w, + S.ple_nq_w, + S.ple_nc_w, + S.ple_conv_w, + conv_states, + si_all, + qsl, + next_norm, + core.eps, + core.hc_count, + _PLE_CLAMP_MIN, + core.ple_dilation, + _PAD_SLOT_ID, + num_tokens, + carry_all, + ) + hc_out[:num_tokens].copy_(stream) + normed_out[:num_tokens].copy_(normed) + hc_out[num_tokens:].zero_() + normed_out[num_tokens:].zero_() + + def _ple_delta(self, hc) -> torch.Tensor: + """A verify step's PLE output for the batch's `n` stream rows `hc`, before it is + added into the stream, through the unfused ops whose intermediates + `SpeculativeRollback` stages: the pre-step pool rows, the id windows and the conv + input.""" + core = self.core + S = self._stack + step = self._step + conv_states, ngram_pool, si_all, carry_all, qsl = step.ple + core.spec_rollback.stage_ple_pools(ngram_pool, conv_states) window = ngram_token_history( - input_ids, ngram_pool, si_all, carry_all, qsl, core.eos_token_id, core.ngram_size + step.input_ids, ngram_pool, si_all, carry_all, qsl, core.eos_token_id, core.ngram_size ) rows = torch.reshape( ngram_hash( @@ -3130,8 +3621,7 @@ def _ple_block(self, hc, out) -> None: key_normed, query_normed, value, core.hc_count, _PLE_CLAMP_MIN ) conv_in = self._grouped_norm(gated, S.ple_nc) - if stage: - core.spec_rollback.stage_ple_rows(window, conv_in) + core.spec_rollback.stage_ple_rows(window, conv_in) # The short conv reads channel-major; a bare transpose is the certified form. convolved = dilated_short_conv_paged( x=torch.transpose(conv_in, 0, 1), @@ -3145,8 +3635,7 @@ def _ple_block(self, hc, out) -> None: silu_activation=True, pad_slot_id=_PAD_SLOT_ID, ) - torch.add(gated, torch.transpose(convolved, 0, 1), out=out[:num_tokens]) - out[num_tokens:].zero_() + return torch.add(gated, torch.transpose(convolved, 0, 1)) def _splice(self, h) -> None: """Break: write the step's vision rows over their placeholder rows of the token @@ -3246,13 +3735,19 @@ def prepare( si_all = mm.state_indices[:batch] hi_all = mm.has_initial_states[:batch] si_p, si_d = torch.split(si_all, [num_prefills, num_decodes], 0) + # The single-token decode kernel takes its slot vector through a 16-byte-aligned + # pointer, and the generation half of the int32 slot vector starts `num_prefills` + # rows in, aligned only when that count is a multiple of four: one copy per step + # makes it aligned for every layer. + si_d = torch.clone(si_d) hi_p = hi_all[:num_prefills] - # `query_start_loc` covers the context prefix; `_long` the whole batch. + # `query_start_loc` covers the context prefix; `_long` the whole batch, and its + # prefix the context rows' chunked recurrence when generation rows run beside them. qsl_all = mm.query_start_loc qsl_p = None if qsl_all is None else qsl_all[: num_prefills + 1] qsl_long = mm.query_start_loc_long + qsl_long_p = qsl_long[: num_prefills + 1] int_idx = None - qsl_long_p = None replay = None if is_verify: # The runtime's arange over the slot pool, from the cache manager (what @@ -3261,8 +3756,6 @@ def prepare( int_idx = core._bound.int_idx[:num_decodes] if core._bound.replay is not None: replay = _replay_work(core, mm, 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, @@ -3372,7 +3865,11 @@ def forward( Gated-DeltaNet layer). Each writes rows `[:n]` of a fixed output allocated here and zeroes rows `[n:]`, so the padding rows stay finite through the row-wise ops and reach neither a state pool nor a real row. With the backend off every step is - eager and unpadded, and the breaks are plain calls.""" + eager and unpadded, and the breaks are plain calls. + + Every hyper-connection site is `qwen4_exp_hc_inject_norm` (the previous block's + injection fused with this site's grouped norm; after a MoE block it also forms that + block's output) and then the mixer (`_hc_mixer`, `_hc_moe_mixer`).""" core = self.core num_tokens = input_ids.shape[0] assert num_tokens == self._step.physical_tokens, ( @@ -3380,10 +3877,8 @@ def forward( f"{self._step.physical_tokens}" ) - # Two views of one `[3, 1, T]` buffer: `pos` (row 0, flat `[T]`: QSA index rope) and - # `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). - pos = torch.reshape(torch.split(position_ids, [1, _QSA_POSITION_AXES - 1], 0)[0], [-1]) + # `pos3`: contiguous `[3, T]` for `use_mrope=True` (the engine hands a slice of a + # larger buffer, whose row stride the ops reject, hence the copy). pos3 = torch.empty( [_QSA_POSITION_AXES, position_ids.shape[-1]], dtype=torch.int32, @@ -3397,17 +3892,17 @@ def forward( dt = h.dtype dev = h.device - # Scratch shared by every layer that uses it; each break rewrites its whole output - # and both packed buffers are fully rewritten by the projection copies. - qkv = torch.empty([num_tokens, core.qkv_width], dtype=dt, device=dev) + # Scratch shared by every layer that uses it; each break rewrites its whole output. + proj_qkv = torch.empty([num_tokens, core.qkv_width], dtype=dt, device=dev) index_scratch = torch.empty([num_tokens, core.index_scratch_width], dtype=dt, device=dev) attn_out = torch.empty([num_tokens, core.q_width], dtype=dt, device=dev) gdn_out = torch.empty( [1, num_tokens, core.num_v_heads, core.head_v_dim], dtype=dt, device=dev ) - ple_out = torch.empty([num_tokens, core.stream], dtype=dt, device=dev) + ple_hc = torch.empty([num_tokens, core.stream], dtype=dt, device=dev) + ple_normed = torch.empty([num_tokens, core.stream], dtype=dt, device=dev) # The CuteDSL routed output, shared by all layers: each layer's unpermute rewrites - # every row, and the post-inject consumes it before the next layer's MoE starts. + # every row, and the next site's injection consumes it before the next layer's MoE. routed_out = ( torch.empty([num_tokens, core.hidden], dtype=dt, device=dev) if core.trunk_moe_backend == "CUTEDSL" @@ -3417,29 +3912,35 @@ def forward( # The embedding repeated across the hc_count residual streams. hc_count = core.hc_count hc = torch.cat([h] * hc_count, -1) - for i, L in enumerate(self._layers): + layers, S = self._layers, self._stack + _, normed = qwen4_exp_hc_inject_norm(hc, None, None, layers[0].ahc_norm, core.eps, hc_count) + for i, L in enumerate(layers): if i == core.ple_layer: - eager_on_graph(hc, ple_out, inner=self._ple_block) - hc = torch.add(hc, ple_out) - mixed, inject = qwen4_exp_gated_residual( - self._grouped_norm(hc, L.ahc_norm), L.ahc_down, L.ahc_up, L.ahc_inject, hc_count - ) + eager_on_graph(hc, ple_hc, ple_normed, inner=self._ple_block) + hc, normed = ple_hc, ple_normed + mixed, inject = self._hc_mixer(normed, L.ahc_first, L.ahc_up, True) if core.is_attn[i]: block_out = self._attention_block( - mixed, i, L, qkv, index_scratch, attn_out, pos, pos3 + mixed, i, L, proj_qkv, index_scratch, attn_out, pos3 ) else: block_out = self._gdn_block(mixed, i, L, gdn_out) - 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, normed = qwen4_exp_hc_inject_norm( + hc, block_out, inject, L.mhc_norm, core.eps, hc_count ) - hc = qwen4_exp_post_inject(hc, self._moe_block(mixed, L, dt, routed_out), inject) + mixed, inject, xq, xsf = self._hc_moe_mixer(normed, L.mhc_first, L.mhc_up, L.moe_gscale) + next_norm = layers[i + 1].ahc_norm if i + 1 < len(layers) else S.fhc_norm + # The MoE block output is formed inside the injection: routed + + # sigmoid(gate) * shared, rounded once. Both [T, hidden] operands are + # released here, not at the next layer's MoE block, whose own peak they + # would otherwise add to. + routed, (gate_logits, shared) = self._moe_block(mixed, xq, xsf, L, routed_out) + hc, normed = qwen4_exp_hc_inject_norm( + hc, routed, inject, next_norm, core.eps, hc_count, shared, gate_logits + ) + del routed, gate_logits, shared # The stack-final mixer with no injection weight is the final norm. - S = self._stack - final, _ = qwen4_exp_gated_residual( - self._grouped_norm(hc, S.fhc_norm), S.fhc_down, S.fhc_up, None, hc_count - ) + final, _ = self._hc_mixer(normed, S.fhc_first, S.fhc_up, False) if not core.mtp_enabled: return final # `lm_head` reads the `[T, hidden]` mix; the MTP module reads the @@ -3919,6 +4420,122 @@ def step_args(self, md: TrtllmAttentionMetadata) -> dict: # --------------------------------------------------------------- blocks --- + def _hc_first(self, normed, first, with_inject): + """The mixer's first half from the site's normed stream: the projection + over `[down; block_inject; 0]` and the SiLU / injection activations. + Returns `(low, inject)`. Between `_HC_TC_FIRST_MIN_ROWS` and + `_HC_TC_FIRST_MAX_ROWS` rows (a per-capture constant) one tcgen05 launch, + `qwen4_exp_hc_decode_tc_first`; otherwise cuBLAS's fp32 logits and + `qwen4_exp_hc_lowrank_act`.""" + core = self.core + if _HC_TC_FIRST_MIN_ROWS <= normed.shape[0] <= _HC_TC_FIRST_MAX_ROWS: + return qwen4_exp_hc_decode_tc_first( + normed, first, core.hc_lowrank, core.hc_count, with_inject + ) + return qwen4_exp_hc_lowrank_act( + cublas_mm(normed, first[:, : core.hc_first_cols], None, torch.float32), + core.hc_lowrank, + core.hc_count, + with_inject, + normed.dtype, + ) + + def _hc_mixer(self, normed, first, up, with_inject): + """One hyper-connection mixer from the site's normed stream: the first + projection over `[down; block_inject; 0]` (fp32 logits), the SiLU and + injection activations, the up projection, and the sigmoid gate with the + stream mean. From `_HC_FUSED_UP_MIN_ROWS` rows (a per-capture constant) + the up projection carries the gate and the mean in its epilogue; below + it cuBLAS's up projection + `qwen4_exp_hc_gate_mean` is faster. + Returns `(mixed, inject)`; `inject` is None without `with_inject` (the + stack-final mixer). + + Neither projection takes the step's `gemm` (`low_m_bf16_mm` under CUDA + graphs): neither FlashInfer low-M kernel writes the first projection's + fp32 logits, and the up projection's K = 320 tiles neither of them.""" + core = self.core + low, inject = self._hc_first(normed, first, with_inject) + if normed.shape[0] >= _HC_FUSED_UP_MIN_ROWS: + return qwen4_exp_hc_up_gate_mean(low, up, normed, core.hc_count), inject + return qwen4_exp_hc_gate_mean(cublas_mm(low, up), normed, core.hc_count), inject + + def _hc_site(self, hc, block_out, inject, norm, first, up, with_inject, normed=None, moe=None): + """One trunk hyper-connection site: the previous block's injection (none at + the stack's first site and after the PLE add), this site's grouped norm and + its mixer. Returns `(hc, mixed, inject)`. The row count is a per-capture + constant: at most `_HC_DECODE_GEMV_MAX_ROWS` rows run two GEMV launches, + larger batches `qwen4_exp_hc_inject_norm` and `_hc_mixer`. `normed`, when + given, is this site's grouped norm of `hc` already computed (the PLE tail's, + with no injection pending); the cuBLAS path starts from it. + + After a MoE block, `block_out` is its routed output and `moe` the + `(gate_logits, shared, topk_weights, topk_rows)` `_moe_block` left + uncombined: the site's first kernel forms the block output itself. With + the top-k pair `block_out` is the runner's permuted FC2 rows and that + kernel (on either route) also forms their finalize. A site handed + `normed` has no pending block, so it never carries `moe`.""" + core = self.core + gate_logits, shared, topk_w, topk_rows = moe if moe is not None else (None,) * 4 + if hc.shape[0] <= _HC_DECODE_GEMV_MAX_ROWS: + n_cols = core.hc_lowrank + (core.hc_count if with_inject else 0) + hc, part, ss = qwen4_exp_hc_decode_first_cute( + hc, + block_out, + inject, + norm, + first, + core.hc_count, + n_cols, + shared, + gate_logits, + topk_w, + topk_rows, + ) + mix = qwen4_exp_hc_decode_mix_cute if hc.shape[0] == 1 else qwen4_exp_hc_decode_mix + mixed, inject = mix(part, ss, up, hc, norm, core.eps, core.hc_count, with_inject) + return hc, mixed, inject + if normed is None: + hc, normed = qwen4_exp_hc_inject_norm( + hc, + block_out, + inject, + norm, + core.eps, + core.hc_count, + shared, + gate_logits, + topk_w, + topk_rows, + ) + mixed, inject = self._hc_mixer(normed, first, up, with_inject) + return hc, mixed, inject + + def _hc_moe_site(self, hc, block_out, inject, norm, first, up, gscale): + """`_hc_site` at an MoE site (always with injection), returning the + routed experts' NVFP4 input (linear scales, the layer's global scale + `gscale`) beside the mix, written by the mixer's own last launch: on the + GEMV route `qwen4_exp_hc_decode_mix_fp4_cute`, from `_HC_FUSED_UP_MIN_ROWS` + rows the fused up projection `qwen4_exp_hc_up_gate_mean_fp4`. Between + the two (only if the thresholds leave a gap) `fp4_quantize` reads the + mix back. Returns `(hc, mixed, inject, xq, xsf)`.""" + core = self.core + if hc.shape[0] <= _HC_DECODE_GEMV_MAX_ROWS: + hc, part, ss = qwen4_exp_hc_decode_first_cute( + hc, block_out, inject, norm, first, core.hc_count, core.hc_lowrank + core.hc_count + ) + mixed, inject, xq, xsf = qwen4_exp_hc_decode_mix_fp4_cute( + part, ss, up, hc, norm, core.eps, core.hc_count, True, gscale + ) + return hc, mixed, inject, xq, xsf + hc, normed = qwen4_exp_hc_inject_norm(hc, block_out, inject, norm, core.eps, core.hc_count) + if normed.shape[0] >= _HC_FUSED_UP_MIN_ROWS: + low, inject = self._hc_first(normed, first, True) + mixed, xq, xsf = qwen4_exp_hc_up_gate_mean_fp4(low, up, normed, core.hc_count, gscale) + return hc, mixed, inject, xq, xsf + mixed, inject = self._hc_mixer(normed, first, up, True) + xq, xsf = fp4_quantize(mixed, gscale, _SF_VEC, False, _SF_LINEAR) + return hc, mixed, inject, xq, xsf + def _grouped_norm(self, x: torch.Tensor, weights: tuple) -> torch.Tensor: """See `PrefillTarget._grouped_norm`.""" core = self.core @@ -3928,118 +4545,169 @@ def _grouped_norm(self, x: torch.Tensor, weights: tuple) -> torch.Tensor: ) def _moe_block( - self, x: torch.Tensor, L: _LayerWeights, dt: torch.dtype, routed_out: torch.Tensor | None - ) -> torch.Tensor: - """See `PrefillTarget._moe_block`.""" + self, x, xq, xsf, L: _LayerWeights, routed_out, gemm=cublas_mm, fold=False + ) -> tuple: + """One MoE layer: NVFP4 routed experts plus a bf16 shared expert, left + uncombined. Returns `(routed, (gate_logits, shared, topk_weights, + topk_rows))`: the next hyper-connection site forms `routed + + sigmoid(gate_logits) * shared` inside its first kernel (`_hc_site`'s + `moe`). With `fold` (trtllm-gen, and the next site takes the combine; a + per-capture constant) the runner also skips its finalize: `routed` is then + its permuted FC2 rows and `topk_weights` / `topk_rows` its routing weights + and index map, which that site's first kernel + (`qwen4_exp_hc_decode_first_cute` or `qwen4_exp_hc_inject_norm`) gathers; + without it both are None and `routed` is the finalized `[T, H]` rows. + + The shared expert and its per-token gate logit run on the core's side + stream, concurrently with the routed path, and join it before this + returns. `x` is the bf16 hidden states (router and shared expert); `xq` / + `xsf` are the same states in NVFP4 with **linear** block scales, from + `_hc_moe_site`. On trtllm-gen the routed path hands the runner the router + logits (its Renormalize routing is bitwise the `renorm_moe_routing_op` + + pre-routed composition); on CuteDSL it is `_cute_dsl_routed`, which writes + `routed_out`. `gemm` is the step's bf16 projection entry (see `forward`); it + launches on the current stream, so the shared expert's GEMMs run on the side + stream.""" core = self.core - # The routing op reads logits from data_ptr() ignoring strides. - logits = cublas_mm(x, L.router) + main = torch.cuda.current_stream() + side = core._side_stream + # Fork: the side stream reads `x` only after everything queued on `main` + # so far. What it allocates is consumed on `main` after the join below, + # and the next layer's fork orders any reuse of those blocks after that. + side.wait_stream(main) + with torch.cuda.stream(side): + shared = gemm(flashinfer_silu_and_mul(gemm(x, L.sh_gu)), L.sh_dn) + gate_logits = gemm(x, L.sh_gate) + # The routing reads its logits as a dense row-major buffer from data_ptr() and + # ignores strides; both projection entries allocate them contiguous. + logits = gemm(x, L.router) if core.trunk_moe_backend == "CUTEDSL": - # fp32 weights: the dtype the unpermute scales by. - topk_ids, topk_w = renorm_moe_routing_op( - logits, topk=core.topk, output_dtype=torch.float32 - ) - xq, xsf = fp4_quantize(x, L.moe_gscale, _SF_VEC, False, _SF_LINEAR) - # The sort pads each expert's (token, expert) pairs to whole tiles and emits - # the tile tables both GEMMs walk. - tile_expert, tile_limit, expanded_to_permuted, permuted_to_expanded, _, live_tiles = ( - moe_sort( - token_selected_experts=topk_ids, - token_final_scales=topk_w, - num_experts=core.num_experts, - top_k=core.topk, - local_expert_offset=0, - local_num_experts=core.num_experts, - tile_tokens_dim=_CUTEDSL_TILE_SIZE, - ) - ) - # FC1 gathers its rows from the unpermuted activation, applies SwiGLU and - # requantizes with the layer's one FC2 input scale. `fp4_quantize` returns the - # linear scales flat; the GEMM reads them as [T, hidden / 16]. - h, h_sf = cute_dsl_nvfp4_gather_grouped_gemm_act_fusion( - input=xq.view(torch.float4_e2m1fn_x2), - weight=L.fc1_w, - input_scale=xsf.view(x.shape[0], core.hidden // _SF_VEC), - weight_scale=L.fc1_s, - alpha=L.alpha1, - tile_idx_to_group_idx=tile_expert, - tile_idx_to_mn_limit=tile_limit, - permuted_idx_to_expanded_idx=permuted_to_expanded, - num_non_exiting_tiles=live_tiles, - global_sf=L.global_sf, - num_experts=core.num_experts, - top_k=core.topk, - num_local_experts=core.num_experts, - local_expert_offset=0, - tile_size=_CUTEDSL_TILE_SIZE, - scaling_vector_size=_SF_VEC, - activation_type=int(ActivationType.Swiglu), - ) - # FC2 writes one bf16 row per permuted row; the unpermute writes every row of - # `routed_out` as the routing-weighted sum of its token's top-k rows. - y = cute_dsl_nvfp4_grouped_gemm( - input=h, - weight=L.fc2_w, - input_scale=h_sf, - weight_scale=L.fc2_s, - alpha=L.alpha2, - tile_idx_to_group_idx=tile_expert, - num_non_exiting_tiles=live_tiles, - num_experts=core.num_experts, - top_k=core.topk, - num_local_experts=core.num_experts, - local_expert_offset=0, - tile_size=_CUTEDSL_TILE_SIZE, - output_dtype=torch.bfloat16, - scaling_vector_size=_SF_VEC, + routed = _cute_dsl_routed(core, L, logits, xq, xsf, routed_out) + main.wait_stream(side) + return routed, (gate_logits, shared, None, None) + moe_op = fp4_block_scale_moe_renorm_permuted if fold else fp4_block_scale_moe_renorm_routed + out = moe_op( + logits, + xq, + xsf.view(torch.float8_e4m3fn), + L.fc1_w, + L.fc1_s, + L.fc2_w, + L.fc2_s, + L.out1_scale, + L.out1_gate_scale, + L.out2_scale, + core.num_experts, + core.topk, + core.moe_inter, + ) + main.wait_stream(side) + if fold: + routed, topk_w, topk_rows = out + return routed, (gate_logits, shared, topk_w, topk_rows) + return out, (gate_logits, shared, None, None) + + def _attention_prologue( + self, proj, q_index, index_view, index_cache, index_section, L, qsa, side, stage + ) -> None: + """The norm + rope of both planes on this stream and the QSA index branch (the + verify step's staging, the raw-key store, then compression) on the side stream, + which reads only the projection's index-key columns and the index pools; the + attention core joins it before reading either.""" + core = self.core + num_tokens = proj.shape[0] + req_idx, logical, _, _, _, _, block_table, tokens_per_block, _, pos3, coords, _ = qsa + side.wait_stream(torch.cuda.current_stream()) + qwen4_exp_qk_norm_rope( + proj, + q_index, + pos3, + core._cos_sin, + L.q_norm, + L.k_norm, + L.idx_qn, + num_heads_q=core.heads_q, + num_heads_k=core.heads_kv, + head_dim=core.head_dim, + index_heads=core.index_heads, + index_head_dim=core.index_head_dim, + index_q_col=core.qkv_width + core.q_width, + rotary_dim=core.rotary_dim, + eps=core.eps, + mrope_section1=core.mrope_section1, + mrope_section2=core.mrope_section2, + index_mrope_section1=index_section[1], + index_mrope_section2=index_section[2], + ) + + # ---- the QSA index branch ------------------------------------------ + # The raw key is stored **before** anything normalizes or rotates it, because + # the compression step pools raw keys and applies the norm and the rotation + # itself at the group's first member's coordinates; a key normalized on the way + # in double-normalizes every anchor that reads it, which is finite, silent, and + # invisible below the sparse threshold. The norm + rope above never writes the + # index key columns. + with torch.cuda.stream(side): + token_k = torch.split( + torch.reshape( + index_view, + [ + num_tokens, + core.index_scratch_width // core.index_head_dim, + core.index_head_dim, + ], + ), + [core.index_heads, core.index_kv_heads, 1], + 1, + )[1] + if stage: + core.spec_rollback.stage_index(index_cache, core._bound.position_cache) + qsa_paged_index_store( + token_k, + coords, + index_cache, + core._bound.position_cache, + req_idx, + logical, + block_table, + tokens_per_block, ) - routed = moe_unpermute_inplace( - permuted_input=y, - output=routed_out, - expanded_idx_to_permuted_idx=expanded_to_permuted, - topk_scales=topk_w, + 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, ) - else: - 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, stage): - """One gated GQA layer plus its QSA index branch for a generation-only batch; - see `PrefillTarget._attention_block`.""" + def _attention_block( + self, x, i, L, qkv, index_scratch, attn_out, qsa, step, stage, gemm=cublas_mm + ): + """One gated GQA layer plus its QSA index branch for a generation-only batch. + + The four input projections are one GEMM against the load-time fused + operand (`weights._attn_in_w`), whose output row is + `[q | k | v | gate | index q, k | v slot]`. One norm+rope launch rewrites + the main q/k heads in place on that output and writes the normalized, + rotated index query into the scratch; the sparse path reads q/k/v as views + of it, the dense path takes one compact copy of its leading `[q | k | v]` + (the dense op reads its QKV flat), the gate is read through its row stride, + and the index branch reads its raw key from it. + + **The index plane is maintained on every step, at both ends of the + threshold.** Below it the dense core runs and the index cache is still + advanced, so a request that grows past the threshold finds its prefix + already compressed; above it the dense core is replaced outright and the + main K/V append becomes this layer's own call, because the op that would + have done it is not the one running. Above the threshold over a bf16 pool + everything up to the index scores is one launch, + `qwen4_exp_attn_decode_prologue`.""" core = self.core ( req_idx, @@ -4057,116 +4725,81 @@ def _attention_block(self, x, i, L, qkv, index_scratch, attn_out, qsa, step, sta ) = 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, - ) - if self._kv_probe is not None and not self._kv_probe.observe(i, k_dst, v_dst): - self._kv_probe = None - - # ---- 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] - if stage: - core.spec_rollback.stage_index(index_cache, core._bound.position_cache) - qsa_paged_index_store( - token_k, - coords, - index_cache, - core._bound.position_cache, - req_idx, - logical, - block_table, - tokens_per_block, - ) - # 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], + proj = gemm(x, L.attn_in) + main = torch.cuda.current_stream() + side = core._side_stream + qkv_view, gate, index_view = torch.split( + proj, [core.qkv_width, core.q_width, core.index_scratch_width], 1 ) + q_dst, k_dst, v_dst = _attn_in_heads(core, proj) + # The index query lands in the q slots of the scratch; its k and v slots + # are layout only and nothing reads them. q_index = torch.split( torch.reshape( - scratch_h, + index_scratch, [num_tokens, core.index_heads + 2 * core.index_kv_heads, core.index_head_dim], ), [core.index_heads, core.index_kv_heads, core.index_kv_heads], 1, )[0] - qsa_prefill_compress( - logical_positions=logical, - request_indices=req_idx, - block_table=block_table, - index_cache=index_cache, - position_cache=core._bound.position_cache, - k_norm_weight=L.idx_kn, - cos_sin=core._cos_sin, - eps=core.eps, - tokens_per_block=tokens_per_block, - compress_ratio=core.compress_ratio, - mrope_section=_QSA_MROPE_SECTION, - ) + index_section = (_QSA_MROPE_SECTION or (0, 0, 0)) if _QSA_INDEX_USE_MROPE else (0, 0, 0) + index_cache = core._bound.index_cache[i] + kv = core._bound.kv[i] + fused_prologue = sparse and kv.k_scale is None and kv.k.dtype == proj.dtype + if fused_prologue: + # The norm + rope of both planes, the K/V append the dense op would have + # owned, and the index branch's raw-key and coordinate store and group + # compression in one launch on this stream. + if stage: + core.spec_rollback.stage_index(index_cache, core._bound.position_cache) + qwen4_exp_attn_decode_prologue( + proj, + q_index, + pos3, + coords, + core._cos_sin, + L.q_norm, + L.k_norm, + L.idx_qn, + L.idx_kn, + kv.k, + kv.v, + index_cache, + core._bound.position_cache, + req_idx, + logical, + block_table, + tokens_per_block, + core.compress_ratio, + num_heads_q=core.heads_q, + num_heads_k=core.heads_kv, + head_dim=core.head_dim, + index_heads=core.index_heads, + index_head_dim=core.index_head_dim, + index_q_col=core.qkv_width + core.q_width, + rotary_dim=core.rotary_dim, + eps=core.eps, + mrope_section1=core.mrope_section1, + mrope_section2=core.mrope_section2, + index_mrope_section1=index_section[1], + index_mrope_section2=index_section[2], + compress_mrope_section=_QSA_MROPE_SECTION, + ) + else: + self._attention_prologue( + proj, q_index, index_view, index_cache, index_section, L, qsa, side, stage + ) + if self._kv_probe is not None and not self._kv_probe.observe(i, k_dst, v_dst): + self._kv_probe = None + gated = False if sparse: - kv = core._bound.kv[i] - # Append the fused-norm+rope'd K/V; the selection reads the current - # token from the pool. Over an fp8 or NVFP4 pool as in the prefill target. - _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block) + if not fused_prologue: + # Selection reads the index plane the side stream wrote. + main.wait_stream(side) + # Append the fused-norm+rope'd K/V; the selection reads the current + # token from the pool. Over an fp8 or NVFP4 pool as in the prefill target. + _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block) scores = qsa_paged_index_scores( q=q_index, index_cache=index_cache, @@ -4177,17 +4810,36 @@ def _attention_block(self, x, i, L, qkv, index_scratch, attn_out, qsa, step, sta compress_ratio=core.compress_ratio, only_visible_blocks=True, context_rows=False, + # Its predecessor is the prologue launch and its consumer the top-k, + # which waits before reading the scores. The prologue writes no + # compressed key but the one the entry re-reads after its wait, so at + # decode-sized batches the pages and keys are read before it. + use_pdl=fused_prologue, + prefetch=fused_prologue and num_tokens <= _SPARSE_DECODE_MAX_ROWS, ) - _select(core, scores, topk_row_starts, visible_blocks, topk_indices) - selected = qsa_expand_block_indices( - topk_indices, - logical, - seq_lens_q, - compress_ratio=core.compress_ratio, - token_topk=core.token_topk, - ) - out = torch.reshape( - _sparse_gqa( + if scores.shape[1] % 4 == 0 and scores.shape[1] <= 4096: + # The canonical block top-k and its expansion to tokens in one launch + # (QSA's row starts are zero). + selected = qsa_topk_expand_decode( + scores, + visible_blocks, + logical, + seq_lens_q, + topk_indices, + core.compress_ratio, + core.token_topk, + ) + else: + _select(core, scores, topk_row_starts, visible_blocks, topk_indices) + selected = qsa_expand_block_indices( + topk_indices, + logical, + seq_lens_q, + compress_ratio=core.compress_ratio, + token_topk=core.token_topk, + ) + if kv.k_scale is not None: + attn = _sparse_gqa( core, kv, q=q_dst, @@ -4198,10 +4850,48 @@ def _attention_block(self, x, i, L, qkv, index_scratch, attn_out, qsa, step, sta # The prefill pair bounds a context row's scan; none here. query_positions=None, compress_ratio=None, - ), - [num_tokens, core.q_width], - ) + ) + else: + # The output gate rides the split-K merge whenever the launch has + # one (more than one split); otherwise it is its own kernel below. + gated = qsa_paged_sparse_gqa.fuses_output_gate(q_dst, kv.k, selected, gate) + if ( + gated + and num_tokens <= _SPARSE_DECODE_MAX_ROWS + and _sparse_decode_supports(q_dst, kv.k, selected) + ): + # The split-K pass, then the merge and the gate as one launch spread + # over the SMs. + attn = qsa_paged_sparse_gqa_decode( + q_dst, + kv.k, + kv.v, + block_table, + selected, + req_idx, + tokens_per_block, + core.softmax_scale, + gate, + ) + else: + attn = qsa_paged_sparse_gqa( + q=q_dst, + k_cache=kv.k, + v_cache=kv.v, + block_table=block_table, + selected_tokens=selected, + request_indices=req_idx, + tokens_per_block=tokens_per_block, + softmax_scale=core.softmax_scale, + query_positions=None, + compress_ratio=None, + output_gate=gate if gated else None, + ) + out = torch.reshape(attn, [num_tokens, core.q_width]) else: + main.wait_stream(side) + # The dense op reads its packed QKV flat: one compact copy. + qkv.copy_(qkv_view) out = thop_attention( q=qkv, output=attn_out, @@ -4215,8 +4905,9 @@ def _attention_block(self, x, i, L, qkv, index_scratch, attn_out, qsa, step, sta **step, **_CALL_CONSTANTS, ) - fused_sigmoid_mul_inplace(out, gate) - return cublas_mm(out, L.o) + if not gated: + fused_sigmoid_mul_inplace(out, gate) + return gemm(out, L.o) def _gdn_verify( self, mq_d, a_d, b_d, ssm_states, i, L, si_d, int_idx, n, p, replay @@ -4283,24 +4974,74 @@ def _gdn_verify( ) return torch.reshape(out, [1, n * p, core.num_v_heads, core.head_v_dim]) - def _gdn_block(self, x, i, L, gdn_state) -> torch.Tensor: + def _gdn_block(self, x, i, L, gdn_state, gemm=cublas_mm, fused_conv=False) -> torch.Tensor: """One Gated-DeltaNet layer over a generation-only batch: one token per - sequence, or under speculation the whole draft chain via the verify path.""" + sequence, or under speculation the whole draft chain through the verify + path. Every branch is on a host int the engine fixes for a capture. + + One GEMM for all four input projections: q|k|v, z, a and b are column views of + its output. A graph-captured single-token step at up to `_GDN_CONV_FUSED_MAX_ROWS` + rows runs the conv step as that GEMM's epilogue (`gdn_in_proj_conv1d_pdl`), and + over an fp32 recurrent pool at up to `_GDN_DECODE_FUSED_MAX_ROWS` rows the + recurrence and the gated norm as one launch (`gdn_decode_gated_norm_pdl`); past + them, after the conv step, as one streaming launch (`gdn_decode_rows_gated_norm`). + Over a bf16 pool the recurrence is `fused_sigmoid_gating_delta_rule_update`'s rows + form, which is batch-invariant.""" core = self.core si_all, qsl_long, is_verify, num_decodes, gen_p, int_idx, replay = gdn_state num_tokens = x.shape[0] - - mixed_qkv = cublas_mm(x, L.in_qkv) - z = cublas_mm(x, L.in_z) - a = cublas_mm(x, L.in_a) - b = cublas_mm(x, L.in_b) - ssm_states = core._bound.ssm[i] + fp32_state = ssm_states.dtype == torch.float32 + + conv_fused = fused_conv and not is_verify and num_tokens <= _GDN_CONV_FUSED_MAX_ROWS + if conv_fused: + proj = gdn_in_proj_conv1d_pdl( + x, + L.in_proj, + core._bound.conv[i], + L.conv_w, + si_all, + core.conv_dim, + _PAD_SLOT_ID, + True, + ) + else: + proj = gemm(x, L.in_proj) + mixed_qkv, z, a, b = torch.split(proj, core.in_proj_split, 1) + + if fp32_state and conv_fused and num_tokens <= _GDN_DECODE_FUSED_MAX_ROWS: + # The recurrence and the gated norm as one launch that reads this layer's + # pool rows before its PDL wait: in a captured single-token step nothing + # ahead of it writes them, and the slot indices are filled before the + # graph is launched. + o = gdn_decode_gated_norm_pdl( + proj, + ssm_states, + si_all, + L.a_log, + L.dt_bias, + L.gnorm, + core.eps, + core.num_k_heads, + core.head_k_dim, + ) + return gemm(o, L.out) if is_verify: out = self._gdn_verify( - mixed_qkv, a, b, ssm_states, i, L, si_all, int_idx, num_decodes, gen_p, replay + mixed_qkv, + a.contiguous(), + b.contiguous(), + ssm_states, + i, + L, + si_all, + int_idx, + num_decodes, + gen_p, + replay, ) - else: + return self._gdn_epilogue(out, z, L, num_tokens, gemm) + if not conv_fused: causal_conv1d_update( x=mixed_qkv, conv_state=core._bound.conv[i], @@ -4311,61 +5052,56 @@ def _gdn_block(self, x, i, L, gdn_state) -> torch.Tensor: conv_state_indices=si_all, pad_slot_id=_PAD_SLOT_ID, ) - if ssm_states.dtype == torch.bfloat16: - # Over a bf16 pool a call with `cu_seqlens` reaches FlashInfer's CuTe-DSL - # decode, whose kernel changes with the batch size and so changes a row's - # last bits; without `cu_seqlens` the entry runs its Triton kernel, one program - # per (row, head). That form finds row n's tokens at `n * q.stride(1)`, so q/k/v - # must be split from the whole contiguous `[T, 10240]` block reshaped to - # `[T, 1, 80, 128]` (K and V heads are both 128 wide), where the size-1 token - # axis carries the row stride; a/b are contiguous `[T, 48]`. - qkv = torch.reshape(mixed_qkv, [num_tokens, 1, -1, core.head_k_dim]) - q, k, v = torch.split( - qkv, [core.num_k_heads, core.num_k_heads, core.num_v_heads], 2 - ) - out = fused_sigmoid_gating_delta_rule_update( - A_log=L.a_log, - a=torch.reshape(a, [num_tokens, 1, core.num_v_heads]), - dt_bias=L.dt_bias, - softplus_beta=_SOFTPLUS_BETA, - softplus_threshold=_SOFTPLUS_THRESHOLD, - q=q, - k=k, - v=v, - b=torch.reshape(b, [num_tokens, 1, core.num_v_heads]), - initial_state_source=ssm_states, - initial_state_indices=si_all, - scale=None, - use_qk_l2norm_in_kernel=True, - cu_seqlens=None, - output=None, - ) - else: - # 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) + if fp32_state: + # The recurrence and the gated norm as one launch over `proj`, whose q/k/v + # columns the conv step rewrote in place; it computes both gates and the + # q/k L2 norm itself, at the softplus pair (1, 20) and epsilon 1e-6 this + # target's constants name, and reads this layer's pool rows before its PDL + # wait (nothing in a single-token step writes them ahead of it). + o = gdn_decode_rows_gated_norm( + proj, + ssm_states, + si_all, + L.a_log, + L.dt_bias, + L.gnorm, + core.eps, + core.num_k_heads, + core.head_k_dim, + ) + return gemm(o, L.out) + else: + # The rows form, `cu_seqlens` None: it finds row n's tokens at + # `n * q.stride(1)`, so q/k/v are split from a contiguous `[T, 10240]` copy + # reshaped to `[T, 1, 80, 128]` (K and V heads are both 128 wide), where the + # size-1 token axis carries the row stride; a/b are contiguous `[T, 48]`. + qkv_rows = torch.reshape(mixed_qkv.contiguous(), [num_tokens, 1, -1, core.head_k_dim]) + q, k, v = torch.split( + qkv_rows, [core.num_k_heads, core.num_k_heads, core.num_v_heads], 2 + ) + out = fused_sigmoid_gating_delta_rule_update( + A_log=L.a_log, + a=torch.reshape(a.contiguous(), [num_tokens, 1, core.num_v_heads]), + dt_bias=L.dt_bias, + softplus_beta=_SOFTPLUS_BETA, + softplus_threshold=_SOFTPLUS_THRESHOLD, + q=q, + k=k, + v=v, + b=torch.reshape(b.contiguous(), [num_tokens, 1, core.num_v_heads]), + initial_state_source=ssm_states, + initial_state_indices=si_all, + scale=None, + use_qk_l2norm_in_kernel=True, + cu_seqlens=None, + output=None, + ) + return self._gdn_epilogue(out, z, L, num_tokens, gemm) - def _gdn_epilogue(self, out, z, L, num_tokens) -> torch.Tensor: - """Gated RMS norm and output projection shared by every recurrence.""" + def _gdn_epilogue(self, out, z, L, num_tokens, gemm=cublas_mm) -> torch.Tensor: + """Gated RMS norm (sigmoid gate, plain non-gemma weight) and output projection + shared by every recurrence.""" core = self.core - # Sigmoid gate, plain (non-gemma) weight. o = rms_norm_gated_token_major( torch.reshape(out, [-1, core.head_v_dim]), torch.reshape(z, [num_tokens, core.num_v_heads, core.head_v_dim]), @@ -4374,58 +5110,109 @@ def _gdn_epilogue(self, out, z, L, num_tokens) -> torch.Tensor: fp8_scale=None, gate_activation="sigmoid", ) - return cublas_mm(torch.reshape(o, [num_tokens, core.value_dim]), L.out) - - def _ple_block(self, hc, input_ids, ple_state, stage) -> torch.Tensor: - """The n-gram side path on one decoder layer; see `PrefillTarget._ple_block`.""" + return gemm(torch.reshape(o, [num_tokens, core.value_dim]), L.out) + + def _ple_block(self, hc, input_ids, ple_state, next_norm, stage, gemm=cublas_mm) -> tuple: + """The n-gram side path on one decoder layer, added into the stream and + normalized for the next site. Returns `(hc, normed)`. + + The same steps as the prefill target's `_ple_block` -- id history, hash, table + gather and dequant, the gated injection, the dilated short conv -- in three + launches plus the two projections, because a decode batch carries at most + `gen_p` rows per sequence: + + 1. `ple_ngram_embed`: window, hash, gather, dequantize, and the id + history published back, one program per sequence; + 2. the key and value projections, on the step's bf16 projection entry + `gemm` (see `forward`); + 3. `ple_gate_conv_inject_norm`: the key / query grouped norms, the gate, + the conv-input grouped norm, the conv against the paged conv state + (published back), the residual add into the stream and this + layer's attention-side grouped norm, one program per + (sequence, stream). + + Every generation sequence carries its history by construction, so no carry + flags are passed. A verify step (`stage`) takes the unfused ops instead, whose + intermediates `SpeculativeRollback` stages (see `PrefillTarget._ple_delta`).""" core = self.core S = self._stack - conv_states, ngram_pool, si_all, carry_all, qsl = ple_state - + conv_states, ngram_pool, si_all, carry_all, qsl, gen_p = ple_state if stage: core.spec_rollback.stage_ple_pools(ngram_pool, conv_states) - window = ngram_token_history( - input_ids, ngram_pool, si_all, carry_all, qsl, core.eos_token_id, core.ngram_size - ) - rows = torch.reshape( - ngram_hash( - window, - layer_multipliers=S.ple_mult, - head_vocab_sizes=S.ple_vocab, - head_offsets=S.ple_offs, - eos_token_id=core.eos_token_id, - num_query_tokens=1, - ), - [-1, core.ngram_heads], - ) - gathered = nn.functional.embedding(rows, core.w["ple_table"]) - ngram = torch.reshape( - dequantize_e4m3_per_tensor(gathered, S.ple_tscale), [-1, core.ple_embed_dim] - ) - - key_normed = self._grouped_norm(cublas_mm(ngram, S.ple_key), S.ple_nk) - query_normed = self._grouped_norm(hc, S.ple_nq) - value = cublas_mm(ngram, S.ple_val) - gated = gated_ngram_injection( - key_normed, query_normed, value, core.hc_count, _PLE_CLAMP_MIN - ) - conv_in = self._grouped_norm(gated, S.ple_nc) - if stage: + window = ngram_token_history( + input_ids, ngram_pool, si_all, carry_all, qsl, core.eos_token_id, core.ngram_size + ) + rows = torch.reshape( + ngram_hash( + window, + layer_multipliers=S.ple_mult, + head_vocab_sizes=S.ple_vocab, + head_offsets=S.ple_offs, + eos_token_id=core.eos_token_id, + num_query_tokens=1, + ), + [-1, core.ngram_heads], + ) + gathered = nn.functional.embedding(rows, core.w["ple_table"]) + ngram = torch.reshape( + dequantize_e4m3_per_tensor(gathered, S.ple_tscale), [-1, core.ple_embed_dim] + ) + key_normed = self._grouped_norm(cublas_mm(ngram, S.ple_key), S.ple_nk) + query_normed = self._grouped_norm(hc, S.ple_nq) + value = cublas_mm(ngram, S.ple_val) + gated = gated_ngram_injection( + key_normed, query_normed, value, core.hc_count, _PLE_CLAMP_MIN + ) + conv_in = self._grouped_norm(gated, S.ple_nc) core.spec_rollback.stage_ple_rows(window, conv_in) - # The short conv reads channel-major; a bare transpose is the certified form. - convolved = dilated_short_conv_paged( - x=torch.transpose(conv_in, 0, 1), - weight=S.ple_conv_w, - bias=None, - conv_states=conv_states, - query_start_loc=qsl, - cache_indices=si_all, - has_initial_state=carry_all, - dilation=core.ple_dilation, - silu_activation=True, - pad_slot_id=_PAD_SLOT_ID, + # The short conv reads channel-major; a bare transpose is the certified form. + convolved = dilated_short_conv_paged( + x=torch.transpose(conv_in, 0, 1), + weight=S.ple_conv_w, + bias=None, + conv_states=conv_states, + query_start_loc=qsl, + cache_indices=si_all, + has_initial_state=carry_all, + dilation=core.ple_dilation, + silu_activation=True, + pad_slot_id=_PAD_SLOT_ID, + ) + hc = torch.add(hc, torch.add(gated, torch.transpose(convolved, 0, 1))) + return qwen4_exp_hc_inject_norm(hc, None, None, next_norm, core.eps, core.hc_count) + ngram = ple_ngram_embed( + input_ids, + ngram_pool, + si_all, + None, + qsl, + core.w["ple_table"], + S.ple_tscale, + S.ple_mult, + S.ple_vocab, + S.ple_offs, + core.eos_token_id, + gen_p, + ) + return ple_gate_conv_inject_norm( + hc, + gemm(ngram, S.ple_key), + gemm(ngram, S.ple_val), + S.ple_nk_w, + S.ple_nq_w, + S.ple_nc_w, + S.ple_conv_w, + conv_states, + si_all, + qsl, + next_norm, + core.eps, + core.hc_count, + _PLE_CLAMP_MIN, + core.ple_dilation, + _PAD_SLOT_ID, + gen_p, ) - return torch.add(gated, torch.transpose(convolved, 0, 1)) # -------------------------------------------------------------- forward --- @@ -4443,6 +5230,9 @@ def forward( step = self.step_args(md) spec_metadata = kwargs.get("spec_metadata") mm = md.mamba_metadata + # The bf16 block projections of a graph-captured step run as PDL launches + # (`_decode_projection`); an eager step keeps cuBLAS. + gemm = _decode_projection if md.is_cuda_graph else cublas_mm # Three views of one `[3, 1, T]` buffer (see PrefillTarget.forward): `pos`, # contiguous `pos3`, and `coords`. @@ -4473,7 +5263,7 @@ def forward( # The batch's bookkeeping, built once; every sequence is a generation sequence. batch = md.num_seqs - # Host-int check (graph-safe); see PrefillTarget.forward. + # Host-int check (graph-safe); see PrefillTarget.prepare. assert num_tokens == md.num_tokens, ( f"the engine handed {num_tokens} activation rows for a batch its " f"metadata describes as {md.num_tokens} tokens " @@ -4497,7 +5287,7 @@ def forward( "speculation off; the base path implements one query token per " "generation sequence" ) - # `gen_p > 1` is load-bearing for verify (see PrefillTarget.forward). + # `gen_p > 1` is load-bearing for verify (see PrefillTarget.prepare). if spec_metadata is not None and batch > 0: declared_p = spec_metadata.runtime_draft_len + 1 assert gen_p == declared_p, ( @@ -4529,12 +5319,19 @@ def forward( [batch], dtype=mm.has_initial_states.dtype, device=mm.has_initial_states.device ) # PLE segmentation vector: rebuilt from `seq_lens_cuda` when `gen_p > 1` (see - # PrefillTarget.forward). + # PrefillTarget.prepare). if gen_p == 1: ple_qsl = qsl_long else: ple_qsl = cu_seqlens_from_seq_lens(md.seq_lens_cuda[:batch], torch.int64) - ple_state = (core._bound.ple_state[0], core._bound.ple_state[1], si_all, ple_carry, ple_qsl) + ple_state = ( + core._bound.ple_state[0], + core._bound.ple_state[1], + si_all, + ple_carry, + ple_qsl, + gen_p, + ) # Dense/sparse decision once per forward, from a host-side metadata read: a # per-capture constant matching the engine's short/long decode graph families. @@ -4556,7 +5353,7 @@ def forward( ) # Staging for the speculative rollback, as on the prefill target (see - # PrefillTarget.forward); every row here is a generation row. + # PrefillTarget.prepare); every row here is a generation row. stage = is_verify rollback = core.spec_rollback if rollback is not None: @@ -4574,12 +5371,14 @@ def forward( tokens_per_block=md.tokens_per_block, ) - # Scratch shared by all attention layers (each fully rewrites them). + # Scratch shared by all attention layers: thop_attention overwrites rows + # [:num_tokens] of `output` in place, and the packed buffers are fully + # rewritten every layer (the QKV one on the dense path only). qkv = torch.empty([num_tokens, core.qkv_width], dtype=dt, device=dev) index_scratch = torch.empty([num_tokens, core.index_scratch_width], dtype=dt, device=dev) attn_out = torch.empty([num_tokens, core.q_width], dtype=dt, device=dev) # The CuteDSL routed output, shared by all layers: each layer's unpermute rewrites - # every row, and the post-inject consumes it before the next layer's MoE starts. + # every row, and the next site's injection consumes it before the next layer's MoE. routed_out = ( torch.empty([num_tokens, core.hidden], dtype=dt, device=dev) if core.trunk_moe_backend == "CUTEDSL" @@ -4589,27 +5388,56 @@ def forward( # The embedding repeated across the hc_count residual streams. hc_count = core.hc_count hc = torch.cat([h] * hc_count, -1) - for i, L in enumerate(self._layers): + layers, S = self._layers, self._stack + # Every hyper-connection site (`_hc_site`): the previous block's pending + # injection, this site's grouped norm and its mixer. The stack's first + # site has no pending injection; at the PLE layer it is applied alone + # (the PLE block reads the injected stream), and the PLE tail returns the + # stream with its output added plus this site's grouped norm of it, which + # the site's cuBLAS mixer starts from (a GEMV site normalizes inside its + # first launch and reads the stream alone). A MoE block's output is + # pending uncombined (`moe`): the next site's first kernel forms it, + # except before the PLE add, where the combine runs alone. + block_out = inject = moe = None + for i, L in enumerate(layers): if i == core.ple_layer: - hc = torch.add(hc, self._ple_block(hc, input_ids, ple_state, stage)) - mixed, inject = qwen4_exp_gated_residual( - self._grouped_norm(hc, L.ahc_norm), L.ahc_down, L.ahc_up, L.ahc_inject, hc_count + if block_out is not None: + if moe is not None: + # `_moe_block` never folds the finalize before this layer. + block_out = fused_sigmoid_gate_mul_add(block_out, *moe[:2]) + hc = qwen4_exp_post_inject(hc, block_out, inject) + block_out = inject = moe = None + hc, normed = self._ple_block(hc, input_ids, ple_state, L.ahc_norm, stage, gemm) + else: + normed = None + hc, mixed, inject = self._hc_site( + hc, + block_out, + inject, + L.ahc_norm, + L.ahc_first, + L.ahc_up, + True, + normed=normed, + moe=moe, ) if core.is_attn[i]: block_out = self._attention_block( - mixed, i, L, qkv, index_scratch, attn_out, qsa, step, stage + mixed, i, L, qkv, index_scratch, attn_out, qsa, step, stage, gemm ) else: - block_out = self._gdn_block(mixed, i, L, gdn_state) - hc = qwen4_exp_post_inject(hc, block_out, inject) - mixed, inject = qwen4_exp_gated_residual( - self._grouped_norm(hc, L.mhc_norm), L.mhc_down, L.mhc_up, L.mhc_inject, hc_count + block_out = self._gdn_block(mixed, i, L, gdn_state, gemm, md.is_cuda_graph) + hc, mixed, inject, xq, xsf = self._hc_moe_site( + hc, block_out, inject, L.mhc_norm, L.mhc_first, L.mhc_up, L.moe_gscale ) - hc = qwen4_exp_post_inject(hc, self._moe_block(mixed, L, dt, routed_out), inject) - # The stack-final mixer with no injection weight is the final norm. - S = self._stack - final, _ = qwen4_exp_gated_residual( - self._grouped_norm(hc, S.fhc_norm), S.fhc_down, S.fhc_up, None, hc_count + # The trtllm-gen runner's finalize folds into the next site's first kernel + # when that site takes the combine (every one but the PLE layer's). + fold = core.trunk_moe_backend == "TRTLLM" and i + 1 != core.ple_layer + block_out, moe = self._moe_block(mixed, xq, xsf, L, routed_out, gemm, fold=fold) + # The stack-final mixer with no injection weight: it returns the mix + # alone, and that mix *is* this checkpoint's final norm. + hc, final, _ = self._hc_site( + hc, block_out, inject, S.fhc_norm, S.fhc_first, S.fhc_up, False, moe=moe ) if spec_metadata is None: return final @@ -5511,6 +6339,33 @@ def commit_speculative_states( } +class _WidenedHead: + """`lm_head` whose bf16 logits come back as fp32 from `bf16_to_fp32` (one vectorized + kernel; torch's `.float()` is an unvectorized copy at under half the bandwidth), so the + logits processor's own `.float()` returns them as they are. The same values bit for + bit.""" + + def __init__(self, lm_head) -> None: + self.lm_head = lm_head + + def __call__(self, hidden_states: torch.Tensor) -> torch.Tensor: + logits = self.lm_head(hidden_states) + if logits.dtype == torch.bfloat16 and logits.is_contiguous(): + return bf16_to_fp32(logits) + return logits + + +class _WideningLogitsProcessor(LogitsProcessor): + """The shell's logits processor (its row gather, its fp32 logits) over `_WidenedHead`.""" + + def forward( + self, hidden_states, lm_head, attn_metadata, return_context_logits: bool = False + ) -> torch.Tensor: + return super().forward( + hidden_states, _WidenedHead(lm_head), attn_metadata, return_context_logits + ) + + class ModelingV2Qwen38FlashNextNvfp4Sm103Tp1( DecoderModelForCausalLM[Qwen4ExpTp1ModelingV2, PretrainedConfig] ): @@ -5552,6 +6407,8 @@ def __init__(self, model_config: ModelConfig): hidden_size=text.hidden_size, vocab_size=text.vocab_size, ) + # The logits' bf16 -> fp32 cast as one vectorized kernel (the draft head shares it). + self.logits_processor = _WideningLogitsProcessor() # Media placeholder ids for the executor; built lazily (`__init__` runs under # meta init). self._mm_token_ids_value = (self.model.image_token_id, self.model.video_token_id) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/weights.py index 0b02e1bfb7e6..83662a561062 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/weights.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/weights.py @@ -9,12 +9,20 @@ Transforms: +* trunk attention input projection: `q_proj`, `k_proj`, `v_proj` and the indexer's + `index_qk_proj` read the same activation, so they are one + `[q | k | v | gate | index | pad, hidden]` operand and one GEMM (`_attn_in_w`). `q_proj` + interleaves q and the output gate per head; its q halves go first so the GEMM output's + leading columns are the packed `[q | k | v]` the fused norm+rope rewrites in place, then + the gate halves, the index projection and `index_head_dim` zero rows (a row is a whole + number of `head_dim` heads). The draft (MTP) layer keeps its separate projections; * fused shared-expert gate|up: `[gate ; up]` (the half order `flashinfer_silu_and_mul` reads); * NVFP4 routed-expert stacks, laid out for the trunk's routed backend (`core.trunk_moe_backend`): `[up ; gate]`, then 128x4 scale swizzle. trtllm-gen also interleaves FC1 row by row and 32-row block shuffles weight and scale bytes (FC2: shuffle only); CuteDSL interleaves FC1 in 64-row up/gate groups and leaves FC2 in - checkpoint row order; + checkpoint row order. Under trtllm-gen a second copy in the CuteDSL layout (`fc*_cd`) + serves the prefill-sized launches (`modeling._CUTE_DSL_MOE_MIN_ROWS`); * n-gram table: 128 row chunks concatenated into one `[rows, 160]` buffer (the hash emits global row indices). @@ -43,7 +51,7 @@ from cuda.bindings import runtime as cudart from torch import nn -from tensorrt_llm._utils import CUASSERT, prefer_pinned +from tensorrt_llm._utils import CUASSERT #: `layers=ALL_LAYERS` is the common case and the default. A sentinel rather #: than a lambda so the table reads as data, and rather than `True` so the @@ -106,7 +114,7 @@ def sources(self, d: SimpleNamespace, i: int | None, prefix: str) -> list[tuple[ # ----------------------------------------------------------- host residency -- # # Allocates the one PLE n-gram table. The table stays a `cuda` tensor of the same -# shape/dtype. Pinned, device-mapped pages (not a host gather) because the gather's +# shape/dtype. Registered, device-mapped pages (not a host gather) because the gather's # row indices are produced on device inside the captured decode graph. #: `cudaMemoryType` values; the core asserts on them after load. @@ -126,10 +134,45 @@ def memory_kind(ptr: int) -> int: return int(attributes.type) -def _host_resident_empty(rows: int, cols: int, dtype, device) -> tuple[torch.Tensor, torch.Tensor]: +#: `cudaHostRegister` flags (`driver_types.h`): mapped, so a kernel may dereference the +#: pages, and portable, for every context in the process. +_CUDA_HOST_REGISTER_PORTABLE = 0x01 +_CUDA_HOST_REGISTER_MAPPED = 0x02 + + +class _RegisteredHostBuffer: + """A pageable CPU allocation registered (pinned and mapped) for the device, and + unregistered before its storage is released: at this size glibc serves the + allocation with `mmap` and returns it with `munmap`, so freeing registered pages + would leave the driver mapping memory the OS may hand out again. + + Not torch's pinned allocator: it rounds the 47.7 GiB table up to 64 GiB, and on + GB300 the device's random row gathers from its pages run an order of magnitude + slower than from registered pageable memory.""" + + def __init__(self, nbytes: int) -> None: + self._storage = torch.empty(nbytes, dtype=torch.uint8) + self.pointer = self._storage.data_ptr() + CUASSERT( + cudart.cudaHostRegister( + self.pointer, nbytes, _CUDA_HOST_REGISTER_PORTABLE | _CUDA_HOST_REGISTER_MAPPED + ) + ) + # Bound now: at interpreter shutdown the module may already be torn down. + self._unregister = cudart.cudaHostUnregister + + def __del__(self) -> None: + pointer, self.pointer = self.pointer, 0 + if pointer: + self._unregister(pointer) + + +def _host_resident_empty( + rows: int, cols: int, dtype, device +) -> tuple[torch.Tensor, _RegisteredHostBuffer]: """`[rows, cols]` of `dtype` on `device`, whose bytes live in host memory. - Returns the tensor and the pinned CPU allocation that owns the pages. The + Returns the tensor and the registered CPU allocation that owns the pages. The caller **must** keep the second alive for as long as the first: the tensor is a view through a device address and owns nothing. """ @@ -140,14 +183,8 @@ def _host_resident_empty(rows: int, cols: int, dtype, device) -> tuple[torch.Ten ) assert device.type == "cuda", f"host offload needs a cuda device to map into; got {device}" - # torch's pinned allocator rounds up to a power of two (47.7 GiB -> 64 GiB resident), - # and keeps the block cached after release. - host = torch.empty(rows * cols, dtype=torch.uint8, pin_memory=prefer_pinned()) - assert host.is_pinned(), ( - "host offload maps pinned memory into the device, which is unavailable " - "here (e.g. Confidential Compute)" - ) - (device_pointer,) = CUASSERT(cudart.cudaHostGetDevicePointer(host.data_ptr(), 0)) + host = _RegisteredHostBuffer(rows * cols) + (device_pointer,) = CUASSERT(cudart.cudaHostGetDevicePointer(host.pointer, 0)) # No public torch API builds a tensor over a pointer the caching allocator # did not hand out; `__cuda_array_interface__` is the numba/cupy path. @@ -190,7 +227,15 @@ def _host_resident_empty(rows: int, cols: int, dtype, device) -> tuple[torch.Ten # relayouts allocate per-expert scratch 512 times per parameter and the table # fill stages one shard at a time; without this, peak load memory is the whole # model rather than one parameter. -_RELEASE_AFTER = ("_fc2_s", "_fc1_s", "ple_table", "mtp_fc1_w", "mtp_fc2_w") +_RELEASE_AFTER = ( + "_fc2_s", + "_fc1_s", + "_fc2_s_cd", + "_fc1_s_cd", + "ple_table", + "mtp_fc1_w", + "mtp_fc2_w", +) def _block32_perm(rows: int, device) -> torch.Tensor: @@ -244,10 +289,33 @@ def _swizzle(x: torch.Tensor) -> torch.Tensor: return torch.ops.trtllm.block_scale_interleave(x.contiguous()) +#: Rows after the shared expert's gate|up rows in its fused operand: the per-token gate's +#: one row, then zeros, so the GEMM's output row stays a whole number of 16-byte vectors. +SH_GATE_ROWS = 8 + + def _cat_rows(t: tuple[torch.Tensor, ...], core) -> torch.Tensor: - """Fused shared-expert gate|up weight bytes: **gate** rows first, the half - order `flashinfer_silu_and_mul` reads.""" - return torch.cat(t, dim=0).contiguous() + """Fused shared-expert weight bytes: **gate** rows first, the half order + `flashinfer_silu_and_mul` reads, then up, then the per-token gate's row and + `SH_GATE_ROWS - 1` zero rows -- one GEMM forms gate|up and the gate logit.""" + gate, up, sh_gate = t + pad = sh_gate.new_zeros(SH_GATE_ROWS - 1, sh_gate.shape[-1]) + return torch.cat((gate, up, sh_gate, pad), dim=0).contiguous() + + +def _attn_in_w(t: tuple[torch.Tensor, ...], core) -> torch.Tensor: + """(q_proj, k_proj, v_proj, index_qk_proj) -> the fused `[q | k | v | gate | index | + pad]` input projection of one trunk attention layer (see the module docstring). + + `q_proj`'s rows are `[heads, 2, head_dim]`: head `h`'s q rows, then its gate rows. + """ + q, k, v, idx = t + heads, head_dim, hidden = core.heads_q, core.head_dim, q.shape[-1] + qg = q.reshape(heads, 2, head_dim, hidden) + pad = q.new_zeros(core.index_head_dim, hidden) + rows = (qg[:, 0].reshape(heads * head_dim, hidden), k, v) + rows += (qg[:, 1].reshape(heads * head_dim, hidden), idx, pad) + return torch.cat(rows, dim=0).contiguous() def _expert_fc1_w(t: tuple[torch.Tensor, ...], core) -> torch.Tensor: @@ -285,6 +353,25 @@ def _expert_fc2_s(t: torch.Tensor, core) -> torch.Tensor: return _swizzle(s).reshape(rows, cols) +#: The layout argument of the second, CuteDSL-layout expert copy (`fc*_cd`). +_CUTEDSL_LAYOUT = SimpleNamespace(trunk_moe_backend="CUTEDSL") + + +def _expert_fc1_w_cd(t: tuple[torch.Tensor, ...], core) -> torch.Tensor: + """(up, gate) packed nibbles -> one expert's FC1 operand in the CuteDSL layout.""" + return _expert_fc1_w(t, _CUTEDSL_LAYOUT) + + +def _expert_fc1_s_cd(t: tuple[torch.Tensor, ...], core) -> torch.Tensor: + """(up, gate) e4m3 scales -> one expert's FC1 scale operand in the CuteDSL layout.""" + return _expert_fc1_s(t, _CUTEDSL_LAYOUT) + + +def _expert_fc2_s_cd(t: torch.Tensor, core) -> torch.Tensor: + """`down_proj` e4m3 scales -> one expert's FC2 scale operand in the CuteDSL layout.""" + return _expert_fc2_s(t, _CUTEDSL_LAYOUT) + + def _mtp_fc1_w(t: tuple[torch.Tensor, ...], core) -> torch.Tensor: """(up, gate) e4m3 codes -> one draft expert's FC1 operand, `[up ; gate]` on the row axis and nothing else. The FP8 block-scale runner reads the @@ -326,6 +413,23 @@ def _gdn_layers(d): return [i for i, a in enumerate(d.is_attn) if not a] +#: The Gated-DeltaNet input projections, in the row order of the concatenated +#: `in_proj_w` (and so the column order of its GEMM output). +_IN_PROJ_PARTS = ("in_proj_qkv", "in_proj_z", "in_proj_a", "in_proj_b") + + +def _in_proj_rows(d, part: int) -> tuple[slice]: + """The row range of `in_proj_w` that input projection `part` fills.""" + start = sum(d.in_proj_split[:part]) + return (slice(start, start + d.in_proj_split[part]),) + + +def _cute_dsl_copy_layers(d): + """Every layer when the trunk's routed experts run on trtllm-gen (the second, CuteDSL + copy that prefill-sized launches read); none under CuteDSL, whose one copy is it.""" + return range(d.num_layers) if d.trunk_moe_backend == "TRTLLM" else () + + def _vision_blocks(d): """The vision tower's blocks. Keyed `l{b}_vis_*`: block `b`, not decoder layer `b`.""" return range(d.v_depth) @@ -457,9 +561,34 @@ def _mixer(tag, prefix, layers): ) -def _gated_attention(tag, src, layers): +def _attn_in_fused(tag, src, layers): + """The trunk attention layers' one input projection: q_proj, k_proj, v_proj and the + QSA indexer's index_qk_proj, re-laid by `_attn_in_w`.""" + keys = ("q_proj.weight", "k_proj.weight", "v_proj.weight", "indexer.index_qk_proj.weight") + return ( + W( + f"{tag}attn_in_w", + shape=lambda d: (d.attn_in_rows, d.hidden), + src=lambda d, i: [(tuple(src(k).format(p=f"{_LM}.layers.{i}") for k in keys), None)], + transform=_attn_in_w, + layers=layers, + ), + ) + + +def _gated_attention(tag, src, layers, fused_in=False): """The gated GQA block's six weights. `q_proj` is twice as wide as q: q and the output - gate interleave per head.""" + gate interleave per head. With `fused_in`, q/k/v are not declared here: they are part + of `_attn_in_fused`'s one operand.""" + return ( + *_attn_in_separate(tag, src, layers, fused_in), + *_attn_out_and_norms(tag, src, layers), + ) + + +def _attn_in_separate(tag, src, layers, fused_in): + if fused_in: + return () return ( W( f"{tag}q_w", @@ -479,6 +608,11 @@ def _gated_attention(tag, src, layers): src=src("v_proj.weight"), layers=layers, ), + ) + + +def _attn_out_and_norms(tag, src, layers): + return ( W( f"{tag}o_w", shape=lambda d: (d.hidden, d.q_width), @@ -490,16 +624,18 @@ def _gated_attention(tag, src, layers): ) -def _index_branch(tag, src, layers): +def _index_branch(tag, src, layers, fused_in=False): """The QSA index branch's three weights: one replicated, unquantized projection - producing `[index_heads q | 1 k]`, and two per-head gemma norms.""" + producing `[index_heads q | 1 k]`, and two per-head gemma norms. With `fused_in`, the + projection is part of `_attn_in_fused`'s one operand and only the norms are declared.""" + projection = W( + f"{tag}idx_w", + shape=lambda d: (d.index_width, d.hidden), + src=src("index_qk_proj.weight"), + layers=layers, + ) return ( - W( - f"{tag}idx_w", - shape=lambda d: (d.index_width, d.hidden), - src=src("index_qk_proj.weight"), - layers=layers, - ), + *(() if fused_in else (projection,)), W( f"{tag}idx_qn", shape=lambda d: (d.index_head_dim,), @@ -516,16 +652,18 @@ def _index_branch(tag, src, layers): def _shared_expert(tag, prefix, layers): - """The bf16 shared expert beside the routed ones, and its per-token gate.""" + """The bf16 shared expert beside the routed ones, its per-token gate in the same + operand as its gate|up (`_cat_rows`).""" return ( W( f"{tag}sh_gu", - shape=lambda d: (2 * d.shared_inter, d.hidden), + shape=lambda d: (2 * d.shared_inter + SH_GATE_ROWS, d.hidden), src=lambda d, i: [ ( ( f"{prefix(i)}.shared_expert.gate_proj.weight", f"{prefix(i)}.shared_expert.up_proj.weight", + f"{prefix(i)}.shared_expert_gate.weight", ), None, ) @@ -539,12 +677,6 @@ def _shared_expert(tag, prefix, layers): src=lambda d, i: [(f"{prefix(i)}.shared_expert.down_proj.weight", None)], layers=layers, ), - W( - f"{tag}sh_gate", - shape=lambda d: (1, d.hidden), - src=lambda d, i: [(f"{prefix(i)}.shared_expert_gate.weight", None)], - layers=layers, - ), ) @@ -570,32 +702,24 @@ class Qwen4ExpTp1Weights: ) ), # --- the gated GQA layers and their QSA index branch --- - *_gated_attention("", lambda s: f"{{p}}.self_attn.{s}", _attn_layers), - *_index_branch("", lambda s: f"{{p}}.self_attn.indexer.{s}", _attn_layers), + # q/k/v and the indexer's replicated, unquantized [index_heads q | 1 k] + # projection as one operand, then o_proj, the q/k norms and the indexer's two + # per-head gemma norms. + *_attn_in_fused("", lambda s: f"{{p}}.self_attn.{s}", _attn_layers), + *_gated_attention("", lambda s: f"{{p}}.self_attn.{s}", _attn_layers, fused_in=True), + *_index_branch("", lambda s: f"{{p}}.self_attn.indexer.{s}", _attn_layers, fused_in=True), # --- the Gated-DeltaNet layers --- - W( - "in_qkv_w", - shape=lambda d: (d.conv_dim, d.hidden), - src="{p}.linear_attn.in_proj_qkv.weight", - layers=_gdn_layers, - ), - W( - "in_z_w", - shape=lambda d: (d.value_dim, d.hidden), - src="{p}.linear_attn.in_proj_z.weight", - layers=_gdn_layers, - ), - # The two per-head scalar projections. - W( - "in_a_w", - shape=lambda d: (d.num_v_heads, d.hidden), - src="{p}.linear_attn.in_proj_a.weight", - layers=_gdn_layers, - ), - W( - "in_b_w", - shape=lambda d: (d.num_v_heads, d.hidden), - src="{p}.linear_attn.in_proj_b.weight", + # The four input projections read the same activation, so they are one + # weight and one GEMM: rows `[q|k|v (conv_dim) | z (value_dim) | a | b]` + # (`d.in_proj_split`), each checkpoint tensor copied into its own row + # range. The forward splits the output into column views. + W( + "in_proj_w", + shape=lambda d: (sum(d.in_proj_split), d.hidden), + src=tuple( + (f"{{p}}.linear_attn.{name}.weight", lambda d, j=j: _in_proj_rows(d, j)) + for j, name in enumerate(_IN_PROJ_PARTS) + ), layers=_gdn_layers, ), W( @@ -660,6 +784,39 @@ class Qwen4ExpTp1Weights: src=_experts(_trunk_mlp, "down_proj.weight_scale"), transform=_expert_fc2_s, ), + # Under trtllm-gen, the same experts again in the CuteDSL layout: prefill-sized + # launches run on them (`modeling._CUTE_DSL_MOE_MIN_ROWS`). + W( + "fc1_w_cd", + shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden // 2), + dtype=u8, + src=_experts(_trunk_mlp, "up_proj.weight", "gate_proj.weight"), + transform=_expert_fc1_w_cd, + layers=_cute_dsl_copy_layers, + ), + W( + "fc1_s_cd", + shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden // SF_VEC), + dtype=u8, + src=_experts(_trunk_mlp, "up_proj.weight_scale", "gate_proj.weight_scale"), + transform=_expert_fc1_s_cd, + layers=_cute_dsl_copy_layers, + ), + W( + "fc2_w_cd", + shape=lambda d: (d.num_experts, d.hidden, d.moe_inter // 2), + dtype=u8, + src=_experts(_trunk_mlp, "down_proj.weight"), + layers=_cute_dsl_copy_layers, + ), + W( + "fc2_s_cd", + shape=lambda d: (d.num_experts, d.hidden, d.moe_inter // SF_VEC), + dtype=u8, + src=_experts(_trunk_mlp, "down_proj.weight_scale"), + transform=_expert_fc2_s_cd, + layers=_cute_dsl_copy_layers, + ), *( W(name, shape=lambda d: (d.num_experts,), dtype=f32, src=_experts(_trunk_mlp, suffix)) for name, suffix in _EXPERT_SCALARS @@ -972,10 +1129,15 @@ def dims(self, core) -> SimpleNamespace: kv_width=cfg.num_key_value_heads * cfg.head_dim, index_head_dim=cfg.indexer_head_dim, index_width=(index_heads + index_kv_heads) * cfg.indexer_head_dim, + # [q | k | v | gate | index q, k | pad of one index head] + attn_in_rows=2 * cfg.num_attention_heads * cfg.head_dim + + 2 * cfg.num_key_value_heads * cfg.head_dim + + (index_heads + index_kv_heads + 1) * cfg.indexer_head_dim, num_v_heads=num_v_heads, head_v_dim=cfg.linear_value_head_dim, conv_dim=2 * key_dim + value_dim, value_dim=value_dim, + in_proj_split=(2 * key_dim + value_dim, value_dim, num_v_heads, num_v_heads), conv_kernel=cfg.linear_conv_kernel_dim, num_experts=cfg.num_experts, moe_inter=cfg.moe_intermediate_size, @@ -995,6 +1157,7 @@ def dims(self, core) -> SimpleNamespace: v_pos_rows=vcfg.num_position_embeddings, v_merge_dim=v_hidden * vcfg.spatial_merge_size * vcfg.spatial_merge_size, mtp_enabled=getattr(core.model_config, "spec_config", None) is not None, + trunk_moe_backend=core.trunk_moe_backend, dtype=core.model_config.torch_dtype, ) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_weight_path.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_weight_path.py index 89d85973e85f..f692c332ac5c 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_weight_path.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_weight_path.py @@ -9,7 +9,7 @@ (`weights._expert_fc1_w` / `_expert_fc1_s` / `_expert_fc2_w` / `_expert_fc2_s`) and scalar derivation (`modeling._moe_layer_scalars`), once per routed backend. The operands then go through the call sequence of the target's `_moe_block` -- CuteDSL: sort, FC1 with SwiGLU, -FC2, unpermute; trtllm-gen: the one-call runner -- with explicit routing. +FC2, unpermute; trtllm-gen: the one-call block-scale runner -- with explicit routing. Two routings hit all 512 experts. In the first every token's only non-zero routing weight is on its own expert, so each expert's output is checked alone: a layout error confined to one @@ -32,6 +32,7 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fp4_block_scale_moe_runner import ( + fp4_block_scale_moe_runner, quantize_nvfp4, ) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 @@ -226,7 +227,7 @@ def _trtllm_gen_routed(layer, x, ids, weights) -> torch.Tensor: """The trtllm-gen branch of `_moe_block`.""" sc = layer.scalars xq, xsf = M.fp4_quantize(x, sc["global_scale"], M._SF_VEC, False, M._SF_LINEAR) - return M.fp4_block_scale_moe_runner( + return fp4_block_scale_moe_runner( routing_logits=None, routing_bias=None, hidden_states=xq, diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py new file mode 100644 index 000000000000..040541c9c4e1 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py @@ -0,0 +1,164 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for ops/moe/fp4_moe_renorm_cute_dsl. + +The experts are the pre-routed entry's test's (drawn NVFP4 codes and +power-of-two scales, so every weight is exact); this entry's layout is built +from their raw per-expert tensors by the target's own `weights.py` transforms, +so the test also pins those. Per cell: inside the routed entry's band against +the fp32 reference, and close to the trtllm-gen routed entry on the same +experts. +""" + +import importlib.util +import zlib +from pathlib import Path + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 + weights as W, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp4_block_scale_moe_renorm_routed import ( # noqa: E501 + fp4_block_scale_moe_renorm_routed as routed, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp4_moe_renorm_cute_dsl import ( # noqa: E501 + fp4_moe_renorm_cute_dsl as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.renorm_moe_routing_op import ( # noqa: E501 + renorm_moe_routing_op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "fp4_moe_renorm_cute_dsl requires a CUDA device" + +DEV = "cuda" +G1 = 137.0 # activation global scale, as the pre-routed entry's test draws it + + +def _helpers(): + path = ( + Path(__file__).resolve().parents[4] + / "moe" + / "test_modeling_v2_fp4_block_scale_moe_runner.py" + ) + spec = importlib.util.spec_from_file_location("_fp4_moe_runner_test_helpers", path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +H = _helpers() +_CD = None + + +def _cute_dsl_experts(ref): + """This entry's operands from the raw expert tensors, through `weights.py`.""" + global _CD + if _CD is None: + up_p, gt_p, up_s, gt_s, dn_p, dn_s = ref["raw"] + e = up_p.shape[0] + fc1_w = torch.stack([W._expert_fc1_w_cd((up_p[i], gt_p[i]), None) for i in range(e)]) + fc1_s = torch.stack([W._expert_fc1_s_cd((up_s[i], gt_s[i]), None) for i in range(e)]) + fc2_w = dn_p.contiguous() + fc2_s = torch.stack([W._expert_fc2_s_cd(dn_s[i], None) for i in range(e)]) + _CD = (fc1_w, fc1_s, fc2_w, fc2_s) + return _CD + + +def _operands(spec, seed): + gen = torch.Generator(device=DEV).manual_seed(seed) + args, ref = H._many_experts() + data, sf, x = H._rand_activation(spec["tokens"], spec["hidden"], G1, gen) + logits = torch.randn(spec["tokens"], spec["num_experts"], generator=gen, device=DEV).to( + spec["logits_dtype"] + ) + ids, _ = renorm_moe_routing_op(logits, spec["top_k"], torch.bfloat16) + g2 = H._calibrate_g2(x, ids, ref) + return args, ref, data, sf, x, logits, g2 + + +def _cute_dsl(ref, data, sf, logits, g2, spec): + fc1_w, fc1_s, fc2_w, fc2_s = _cute_dsl_experts(ref) + s1, sg, s2 = H._scalars(spec["num_experts"], G1, g2) + gsf = torch.tensor([g2], dtype=torch.float32, device=DEV) + return op( + logits, + data, + sf, + fc1_w, + fc1_s, + sg, + fc2_w, + fc2_s, + s2, + gsf, + spec["num_experts"], + spec["top_k"], + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + spec = cell.spec + seed = zlib.crc32(cell.why.encode()) % 2**31 + args, ref, data, sf, x, logits, g2 = _operands(spec, seed) + y = _cute_dsl(ref, data, sf, logits, g2, spec) + assert y.shape == (spec["tokens"], spec["hidden"]) and y.dtype == torch.bfloat16 + exp = op.reference(logits, x, spec["top_k"], ref["up"], ref["gate"], ref["down"], g2) + H._compare(y, exp.to(torch.bfloat16)) + + +@pytest.mark.parametrize("tokens", [1024, 4096]) +def test_close_to_the_trtllm_gen_entry(tokens): + """Same experts, same routing: within the routed entry's own band of each + other's output (the two differ only in accumulation order).""" + spec = next(c.spec for c in op.CELLS if c.spec["tokens"] == 1024) + spec = {**spec, "tokens": tokens} + args, ref, data, sf, x, logits, g2 = _operands(spec, 4000 + tokens) + y = _cute_dsl(ref, data, sf, logits, g2, spec) + s1, sg, s2 = H._scalars(spec["num_experts"], G1, g2) + y_gen = routed( + logits, + data, + sf.view(torch.float8_e4m3fn), + args["gemm1_weights"], + args["gemm1_weights_scale"], + args["gemm2_weights"], + args["gemm2_weights_scale"], + s1, + sg, + s2, + spec["num_experts"], + spec["top_k"], + spec["intermediate"], + ) + rel = ((y.float() - y_gen.float()).norm() / y_gen.float().norm()).item() + print(f"\n {tokens} rows: relative RMS difference from trtllm-gen {rel:.2e}") + assert rel < 1e-2 + + +def test_guard_refuses_a_per_expert_fc2_scale(): + spec = next(c.spec for c in op.CELLS if c.spec["tokens"] == 1024) + args, ref, data, sf, x, logits, g2 = _operands(spec, 9) + fc1_w, fc1_s, fc2_w, fc2_s = _cute_dsl_experts(ref) + s1, sg, s2 = H._scalars(spec["num_experts"], G1, g2) + with validating(op), pytest.raises(AssertionError): + op( + logits, + data, + sf, + fc1_w, + fc1_s, + sg, + fc2_w, + fc2_s, + s2, + s2, + spec["num_experts"], + spec["top_k"], + ) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_ngram_embed.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_ngram_embed.py new file mode 100644 index 000000000000..ea2c43621f6c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_ngram_embed.py @@ -0,0 +1,223 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the ple_ngram_embed entry. + +Integer bookkeeping plus a one-rounding dequantize, so gated bitwise -- the +embedding and the whole published pool -- against the entry's reference, which +is the scalar references of `ngram_token_history` and `ngram_hash` and a +float64 gather-and-scale. Beyond the cells: bitwise agreement with the four +calls it replaces, and graph capture. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.embedding.ngram_hash import ( # noqa: E501 + CKPT_HEAD_OFFSETS, + CKPT_HEAD_VOCAB_SIZES, + CKPT_MULTIPLIERS, + ngram_hash, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.ngram_token_history import ( # noqa: E501 + ngram_token_history, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.ple_ngram_embed import ( # noqa: E501 + CKPT_EOS_TOKEN_ID, + CKPT_HEAD_DIM, + CKPT_SCALE, + CKPT_TABLE_ROWS, + ENGINE_SLOTS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.ple_ngram_embed import ( # noqa: E501 + ple_ngram_embed as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.quantization.dequantize_e4m3_per_tensor import ( # noqa: E501 + dequantize_e4m3_per_tensor, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.weights import ( # noqa: E501 + _host_resident_empty, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "ple_ngram_embed requires a CUDA device" + +DEV = "cuda" +EOS = CKPT_EOS_TOKEN_ID +VOCAB = 248320 +FP8 = torch.float8_e4m3fn +#: Keeps the hosts alive for as long as the tables they back. +_HOSTS = [] + + +def _scattered(n: int, pool_slots: int, g) -> list: + """`n` distinct slots over the whole pool, both ends included.""" + middle = (torch.randperm(pool_slots - 2, generator=g) + 1).tolist() + slots = ([0, pool_slots - 1] + middle)[:n] + assert len(set(slots)) == n and max(slots) < pool_slots + return slots + + +def _geometry(spec, g): + """`(multipliers, vocab, offsets)` for the cell: the checkpoint's own for a + 3-gram of 8 heads per order, deterministic stand-ins otherwise.""" + n = spec.get("ngram_size", 3) + hp = spec.get("heads_per_ngram", 8) + if n == 3 and hp == 8 and spec.get("table") == "ckpt": + return CKPT_MULTIPLIERS, CKPT_HEAD_VOCAB_SIZES, CKPT_HEAD_OFFSETS + mult = ( + CKPT_MULTIPLIERS + if n == 3 + else tuple(int(x) | 1 for x in torch.randint(10**12, 3 * 10**13, (n,), generator=g)) + ) + vocab = tuple( + 1000 + 7 * h + int(x) + for h, x in enumerate(torch.randint(0, 50, (hp * (n - 1),), generator=g)) + ) + offsets, acc = [], 0 + for v in vocab: + offsets.append(acc) + acc += v + return mult, vocab, tuple(offsets) + + +def _table(kind: str, rows: int, touched: torch.Tensor, g) -> torch.Tensor: + if kind == "ckpt": + table = torch.empty(CKPT_TABLE_ROWS, CKPT_HEAD_DIM, dtype=FP8, device=DEV) + table[touched] = torch.randn(touched.numel(), CKPT_HEAD_DIM, generator=g).to(FP8).to(DEV) + return table + src = torch.randn(rows, CKPT_HEAD_DIM, generator=g).to(FP8) + if kind == "host": + table, host = _host_resident_empty(rows, CKPT_HEAD_DIM, FP8, torch.device(DEV)) + _HOSTS.append(host) + table.copy_(src.to(DEV)) + return table + return src.to(DEV) + + +def _with_eos(x: torch.Tensor, density: float, g) -> torch.Tensor: + if density > 0.0: + x = torch.where(torch.rand(x.shape, generator=g) < density, torch.full_like(x, EOS), x) + return x + + +def _build(spec, seed): + g = torch.Generator().manual_seed(seed) + lengths = spec["lengths"] + n = spec.get("ngram_size", 3) + num_seqs = len(lengths) + pool_slots = spec.get("pool_slots", num_seqs + 2) + density = spec.get("eos_density", 0.0) + ids = _with_eos(torch.randint(0, VOCAB, (sum(lengths),), generator=g), density, g) + ids = ids.to(spec.get("ids_dtype", torch.int32)).to(DEV) + pool = _with_eos(torch.randint(0, VOCAB, (pool_slots, n - 1), generator=g), density, g).to(DEV) + slots = _scattered(num_seqs, pool_slots, g) if pool_slots >= 64 else list(range(num_seqs)) + slots = torch.tensor(slots, dtype=torch.int32, device=DEV) + carries = spec.get("carries") + carries = None if carries is None else torch.tensor(carries, dtype=torch.bool, device=DEV) + starts = [0] + for length in lengths: + starts.append(starts[-1] + length) + qsl = torch.tensor(starts, dtype=torch.int64, device=DEV) + mult, vocab, offsets = _geometry(spec, g) + mult = torch.tensor(mult, dtype=torch.int64, device=DEV) + vocab = torch.tensor(vocab, dtype=torch.int64, device=DEV) + offsets = torch.tensor(offsets, dtype=torch.int64, device=DEV) + # The rows this batch touches, from the two scalar references, so that the + # full-height table needs bytes only where they are read. + seeded = torch.ones(num_seqs, dtype=torch.bool, device=DEV) if carries is None else carries + window, _ = ngram_token_history.reference(ids, pool, slots, seeded, qsl, EOS, n) + touched = ngram_hash.reference(window, mult, vocab, offsets, EOS, 1).reshape(-1) + rows = int(offsets[-1] + vocab[-1]) + table = _table(spec.get("table", "small"), rows, touched, g) + scale = torch.tensor([CKPT_SCALE], dtype=torch.bfloat16, device=DEV).reshape(1, 1, 1) + max_rows = spec.get("max_rows", max(lengths)) + return [ids, pool, slots, carries, qsl, table, scale, mult, vocab, offsets, EOS, max_rows] + + +def _compare(out, ref) -> None: + (emb, pool), (ref_emb, ref_pool) = out, ref + assert emb.dtype == ref_emb.dtype and emb.shape == ref_emb.shape + assert emb.is_contiguous() + assert torch.equal(emb, ref_emb), ( + f"embedding differs at {(emb != ref_emb).any(-1).nonzero().flatten()[:8].tolist()}" + ) + assert torch.equal(pool, ref_pool), "published pool differs" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) + ref = op.reference(*args) + pool_before = args[1].clone() + out = op(*args) + _compare((out, args[1]), ref) + if cell.spec.get("table") == "host": + assert args[1].data_ptr() != 0 and not torch.equal(args[1], pool_before) + + +def _old(ids, pool, slots, carries, qsl, table, scale, mult, vocab, offsets, eos, max_rows): + """The four calls the decode PLE block made before this entry.""" + seeded = torch.ones_like(slots, dtype=torch.bool) if carries is None else carries + window = ngram_token_history(ids, pool, slots, seeded, qsl, eos, mult.numel()) + rows = torch.reshape(ngram_hash(window, mult, vocab, offsets, eos, 1), [-1, vocab.numel()]) + gathered = torch.nn.functional.embedding(rows, table) + return torch.reshape(dequantize_e4m3_per_tensor(gathered, scale), [rows.shape[0], -1]) + + +@pytest.mark.parametrize( + "spec", + [ + dict(lengths=(1,) * 128, pool_slots=ENGINE_SLOTS, eos_density=0.05), + dict(lengths=(4,) * 32, pool_slots=ENGINE_SLOTS, eos_density=0.05), + dict(lengths=(1,) * 8, table="host"), + dict(lengths=(8192,), carries=(False,), pool_slots=ENGINE_SLOTS, eos_density=0.01), + dict( + lengths=(5000, 3000) + (1,) * 32, + carries=(False, True) + (True,) * 32, + pool_slots=ENGINE_SLOTS, + eos_density=0.05, + ), + ], + ids=["decode128", "mtp4x32", "host8", "prefill8192", "mixed2ctx32dec"], +) +def test_matches_the_calls_it_replaces(spec) -> None: + args = _build(spec, seed=len(spec["lengths"])) + old_args = list(args) + old_args[1] = args[1].clone() + want = _old(*old_args) + got = op(*args) + assert torch.equal(got, want) + assert torch.equal(args[1], old_args[1]) + + +def test_guard_refuses_short_query_start_loc() -> None: + args = _build(dict(lengths=(1, 1, 1)), seed=3) + args[4] = args[4][:-1] + with validating(op), pytest.raises(AssertionError): + op(*args) + + +def test_cuda_graph_capture_and_replay() -> None: + args = _build(dict(lengths=(1,) * 16, pool_slots=ENGINE_SLOTS, eos_density=0.1), seed=11) + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + op(*args) + torch.cuda.current_stream().wait_stream(warmup) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_out = op(*args) + fresh = _build(dict(lengths=(1,) * 16, pool_slots=ENGINE_SLOTS, eos_density=0.1), seed=12) + args[0].copy_(fresh[0]) + pool_before = args[1].clone() + want, want_pool = op.reference(*args) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(graph_out, want) + assert torch.equal(args[1], want_pool) + assert not torch.equal(args[1], pool_before) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_block.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_block.py new file mode 100644 index 000000000000..418852a9b93d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_block.py @@ -0,0 +1,407 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test: the target's whole Gated-DeltaNet block against the all-Triton path. + +Drives `DecodeTarget._gdn_block` and `PrefillTarget._gdn_block` at the +checkpoint's geometry (hidden 2560, 16 k-heads, 48 v-heads, head dim 128, +fp32 recurrent pool of 461 slots) on decode, prefill-only and mixed +(context + one-token generation rows) batches, and compares the block output +and both state pools with a reference that composes the Triton recurrences +(`fused_sigmoid_gating_delta_rule_update` for decode, +`chunk_gated_delta_rule` over every row of a prefill step) around the same +projections, convolutions and epilogue. + +The reference's context rows run the unfused `extract_transpose_prefill_slice` +-> `causal_conv1d_fwd` -> `fused_gdn_post_conv` chain the target's +`gdn_prefill_conv_post` replaces. + +Bands: the conv pools are bitwise (both sides publish the raw projection). The +recurrent pool is held to the looser of the two recurrence entries' state +bands (`flashinfer_chunk_gated_delta_rule`: rtol 1.6e-2, atol 2e-2), and the +block output, after the gated norm and the output projection, to a relative +L2 distance of 1e-2 and a max-abs distance of 3e-2 of its own peak. +""" + +import types + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 + modeling as M, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.causal_conv1d_fwd import ( # noqa: E501 + causal_conv1d_fwd, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.chunk_gated_delta_rule import ( # noqa: E501 + chunk_gated_delta_rule as triton_chunk, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.extract_transpose_prefill_slice import ( # noqa: E501 + extract_transpose_prefill_slice, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_gdn_post_conv import ( # noqa: E501 + fused_gdn_post_conv, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_sigmoid_gating_delta_rule_update import ( # noqa: E501 + fused_sigmoid_gating_delta_rule_update as triton_decode, +) + +assert torch.cuda.is_available(), "the Gated-DeltaNet block requires a CUDA device" + +DEV = "cuda" +BF16 = torch.bfloat16 +HIDDEN, H, HV, K, V, WIDTH = 2560, 16, 48, 128, 128, 4 +KD, VD = H * K, HV * V +CONV = 2 * KD + VD +SLOTS = 461 +MAX_BATCH = 256 +RTOL_STATE, ATOL_STATE = 1.6e-2, 2e-2 +OUT_REL_L2, OUT_MAX_FRAC = 1e-2, 3e-2 + + +def _slots(n: int) -> list[int]: + """`n` distinct slots of the pool, both ends first.""" + rest = [(i * 37 + 4) % (SLOTS - 2) + 1 for i in range(max(n - 2, 0))] + return ([SLOTS - 1, 0] + rest)[:n] + + +def _core(seed: int) -> types.SimpleNamespace: + """The attributes `_gdn_block` reads off the core, with layer 0's pools.""" + g = torch.Generator(device=DEV).manual_seed(seed) + conv = torch.randn(SLOTS, CONV, WIDTH - 1, generator=g, device=DEV, dtype=BF16) + ssm = torch.randn(SLOTS, HV, V, K, generator=g, device=DEV) * 0.05 + return types.SimpleNamespace( + conv_dim=CONV, + key_dim=KD, + value_dim=VD, + num_k_heads=H, + head_k_dim=K, + num_v_heads=HV, + head_v_dim=V, + eps=1e-6, + conv_kernel=WIDTH, + in_proj_split=(CONV, VD, HV, HV), + _bound=types.SimpleNamespace(conv={0: conv}, ssm={0: ssm}), + ) + + +def _clone_core(core: types.SimpleNamespace) -> types.SimpleNamespace: + out = types.SimpleNamespace(**vars(core)) + out._bound = types.SimpleNamespace( + conv={0: core._bound.conv[0].clone()}, ssm={0: core._bound.ssm[0].clone()} + ) + return out + + +def _layer(seed: int) -> types.SimpleNamespace: + """Layer weights as the target views them: HF `[out, in]` storage read + through `.t()`, fp32 gating vectors over the checkpoint's ranges.""" + g = torch.Generator(device=DEV).manual_seed(seed) + + def lin(n_out, n_in): + return (torch.randn(n_out, n_in, generator=g, device=DEV) * n_in**-0.5).to(BF16).t() + + # The target's one [q|k|v | z | a | b] input projection, read through `.t()`. + # The reference blocks run the same GEMM, so the conv pools stay bitwise. + return types.SimpleNamespace( + in_proj=lin(CONV + VD + 2 * HV, HIDDEN), + out=lin(HIDDEN, VD), + conv_w=(torch.randn(CONV, WIDTH, generator=g, device=DEV) * 0.5).to(BF16), + a_log=torch.empty(HV, device=DEV).uniform_(-5.5, 3.5, generator=g), + dt_bias=torch.empty(HV, device=DEV).uniform_(-8.9, 19.2, generator=g), + gnorm=(1.0 + 0.1 * torch.randn(VD // HV, generator=g, device=DEV)).to(BF16), + ) + + +def _x(tokens: int, seed: int) -> torch.Tensor: + g = torch.Generator(device=DEV).manual_seed(seed) + return torch.randn(tokens, HIDDEN, generator=g, device=DEV, dtype=BF16) + + +def _target(cls, core, L=None, step=None): + """A target over `core` alone; a prefill target's `_gdn_core` break reads layer 0's + weights and the step's bookkeeping off it.""" + target = object.__new__(cls) + target.core = core + target._layers = [L] + target._step = step + return target + + +def _prefill_block(core, x, L, st): + """`PrefillTarget._gdn_block` over the step `st`, the break's output buffer allocated + as the forward allocates it.""" + tokens = x.shape[0] + step = types.SimpleNamespace(gdn=st, num_tokens=tokens) + out = torch.empty([1, tokens, HV, V], dtype=BF16, device=DEV) + return _target(M.PrefillTarget, core, L, step)._gdn_block(x, 0, L, out) + + +def _decode_state(m: int): + """`DecodeTarget.forward`'s `gdn_state` for `m` generation sequences: the + slot vector is the head of an engine-sized buffer.""" + buf = torch.zeros(MAX_BATCH, dtype=torch.int32, device=DEV) + buf[:m] = torch.tensor(_slots(m), dtype=torch.int32, device=DEV) + si_all = buf[:m] + qsl_long = torch.arange(m + 1, dtype=torch.int64, device=DEV) + return (si_all, qsl_long, False, m, 1, None, None), qsl_long + + +def _prefill_state(ctx_lens: tuple, n_dec: int, has_init: tuple): + """`PrefillTarget.forward`'s `gdn_state` for a context prefix plus `n_dec` + one-token generation rows, built the way the forward builds it from the + engine's `Mamba2Metadata` buffers.""" + n_ctx = len(ctx_lens) + batch = n_ctx + n_dec + n_p = sum(ctx_lens) + si_buf = torch.zeros(MAX_BATCH, dtype=torch.int32, device=DEV) + si_buf[:batch] = torch.tensor(_slots(batch), dtype=torch.int32, device=DEV) + hi_buf = torch.zeros(MAX_BATCH, dtype=torch.bool, device=DEV) + hi_buf[:n_ctx] = torch.tensor(has_init, dtype=torch.bool, device=DEV) + si_all, hi_all = si_buf[:batch], hi_buf[:batch] + si_p, si_d = torch.split(si_all, [n_ctx, n_dec], 0) + si_d = torch.clone(si_d) + hi_p = hi_all[:n_ctx] + bounds = [0] + for n in ctx_lens: + bounds.append(bounds[-1] + n) + bounds += [n_p + j for j in range(1, n_dec + 1)] + cu_buf = torch.zeros(MAX_BATCH + 1, dtype=torch.int32, device=DEV) + cu_buf[: batch + 1] = torch.tensor(bounds, dtype=torch.int32, device=DEV) + cu_long_buf = cu_buf.to(torch.int64) + qsl_all = cu_buf[: batch + 1] + qsl_p = qsl_all[: n_ctx + 1] + qsl_long = cu_long_buf[: batch + 1] + qsl_long_p = qsl_long[: n_ctx + 1] + return ( + n_p, + n_dec, + si_p, + si_all, + si_d, + hi_p, + hi_all, + qsl_p, + qsl_all, + qsl_long, + False, + n_dec, + 1, + None, + qsl_long_p, + None, + ) + + +def _post_conv(prefill, decode, a, b, L): + return fused_gdn_post_conv( + prefill=prefill, + decode=decode, + a=a, + b=b, + A_log=L.a_log, + dt_bias=L.dt_bias, + num_k_heads=H, + head_k_dim=K, + num_v_heads=HV, + head_v_dim=V, + l2_norm_eps=M._L2_NORM_EPS, + softplus_threshold=M._SOFTPLUS_THRESHOLD, + ) + + +def _triton_decode_block(core, x, L, si_all, qsl_long): + """The decode block with the Triton single-token recurrence.""" + m = x.shape[0] + mixed, z, a, b = torch.split(M.cublas_mm(x, L.in_proj), [CONV, VD, HV, HV], 1) + M.causal_conv1d_update( + x=mixed, + conv_state=core._bound.conv[0], + weight=L.conv_w, + bias=None, + silu_activation=True, + cache_seqlens=None, + conv_state_indices=si_all, + pad_slot_id=M._PAD_SLOT_ID, + ) + qv, kv, vv = torch.split(mixed, [KD, KD, VD], 1) + out = triton_decode( + A_log=L.a_log, + a=a, + dt_bias=L.dt_bias, + softplus_beta=M._SOFTPLUS_BETA, + softplus_threshold=M._SOFTPLUS_THRESHOLD, + q=torch.reshape(qv, [1, m, H, K]), + k=torch.reshape(kv, [1, m, H, K]), + v=torch.reshape(vv, [1, m, HV, V]), + b=b, + initial_state_source=core._bound.ssm[0], + initial_state_indices=si_all, + scale=None, + use_qk_l2norm_in_kernel=True, + cu_seqlens=qsl_long, + output=None, + ) + return M.DecodeTarget._gdn_epilogue(_target(M.DecodeTarget, core), out, z, L, m) + + +def _triton_prefill_block(core, x, L, st): + """The prefill block with every row -- context and generation -- in one + Triton chunk call.""" + n_p, n_d, si_p, si_all, si_d, hi_p, hi_all, qsl_p, qsl_all, qsl_long = st[:10] + T = x.shape[0] + conv, ssm = core._bound.conv[0], core._bound.ssm[0] + mixed, z, a, b = torch.split(M.cublas_mm(x, L.in_proj), [CONV, VD, HV, HV], 1) + M.reset_recurrent_state_rows(ssm, si_p, hi_p, conv) + conv_args = dict( + weight=L.conv_w, + bias=None, + conv_states=conv, + silu_activation=True, + pad_slot_id=M._PAD_SLOT_ID, + out=None, + ) + if n_d > 0: + mq_p, mq_d = torch.split(mixed, [n_p, n_d], 0) + mq_p_t = extract_transpose_prefill_slice(mq_p, n_p, 0, CONV) + causal_conv1d_fwd( + x=mq_p_t, query_start_loc=qsl_p, cache_indices=si_p, has_initial_state=hi_p, **conv_args + ) + M.causal_conv1d_update( + x=mq_d, + conv_state=conv, + weight=L.conv_w, + bias=None, + silu_activation=True, + cache_seqlens=None, + conv_state_indices=si_d, + pad_slot_id=M._PAD_SLOT_ID, + ) + q, k, v, g, beta = _post_conv(mq_p_t, mq_d, a, b, L) + else: + mq_t = extract_transpose_prefill_slice(mixed, T, 0, CONV) + causal_conv1d_fwd( + x=mq_t, + query_start_loc=qsl_all, + cache_indices=si_all, + has_initial_state=hi_all, + **conv_args, + ) + q, k, v, g, beta = _post_conv(mq_t, None, a, b, L) + out, _ = triton_chunk( + q, + k, + v, + g, + beta, + scale=None, + initial_state=ssm, + initial_state_indices=si_all, + inplace_indexed_state_update=True, + output_final_state=False, + cu_seqlens=qsl_long, + use_qk_l2norm_in_kernel=False, + output=None, + ) + return M.PrefillTarget._gdn_epilogue(_target(M.PrefillTarget, core), out, z, L, T) + + +def _check(y, y_ref, core, core_ref, what): + """Output and both pools against the reference block's, at the module bands.""" + assert y.shape == y_ref.shape and y.dtype == y_ref.dtype + assert torch.isfinite(y).all() + d = y.float() - y_ref.float() + rel = (d.norm() / y_ref.float().norm()).item() + peak = (d.abs().max() / y_ref.float().abs().max()).item() + s_d = core._bound.ssm[0] - core_ref._bound.ssm[0] + s_rel = (s_d.norm() / core_ref._bound.ssm[0].norm()).item() + print( + f"{what}: out rel-L2 {rel:.3e}, out max-abs/peak {peak:.3e}, " + f"state rel-L2 {s_rel:.3e}, state max-abs {s_d.abs().max().item():.3e}" + ) + assert rel <= OUT_REL_L2 and peak <= OUT_MAX_FRAC, (rel, peak) + assert torch.equal(core._bound.conv[0], core_ref._bound.conv[0]), "conv pools differ" + torch.testing.assert_close( + core._bound.ssm[0], core_ref._bound.ssm[0], rtol=RTOL_STATE, atol=ATOL_STATE + ) + + +@pytest.fixture(scope="module") +def warm_triton_chunk(): + """The Triton chunk op's first indexed in-place call per autotune key + returns a wrong answer (its entry's note); spend it on scratch.""" + core = _core(0) + L = _layer(0) + st = _prefill_state((64,), 0, (False,)) + _triton_prefill_block(core, _x(64, 0), L, st) + torch.cuda.synchronize() + + +@pytest.mark.parametrize("m", (1, 32, 128)) +def test_decode_block(m) -> None: + core, L, x = _core(100 + m), _layer(200 + m), _x(m, 300 + m) + core_ref = _clone_core(core) + st, qsl_long = _decode_state(m) + y = _target(M.DecodeTarget, core)._gdn_block(x, 0, L, st) + y_ref = _triton_decode_block(core_ref, x, L, st[0], qsl_long) + _check(y, y_ref, core, core_ref, f"decode m={m}") + named = set(_slots(m)) + unnamed = [s for s in range(SLOTS) if s not in named] + original = _core(100 + m) + assert torch.equal(core._bound.ssm[0][unnamed], original._bound.ssm[0][unnamed]), ( + "unnamed slots written" + ) + + +#: (context lengths, generation rows, has_initial_state per context). +PREFILL_CASES = ( + ((8192,), 0, (False,)), + ((200, 37, 91), 0, (True, False, True)), + ((8065,), 127, (True,)), + ((200, 37, 91), 5, (True, False, True)), + ((64,), 1, (False,)), +) + + +@pytest.mark.parametrize( + "case", PREFILL_CASES, ids=lambda c: f"{'+'.join(map(str, c[0]))}+{c[1]}x1" +) +def test_prefill_block(case, warm_triton_chunk) -> None: # noqa: F811 — fixture by name + ctx_lens, n_dec, has_init = case + tokens = sum(ctx_lens) + n_dec + core, L, x = _core(400 + tokens), _layer(500 + tokens), _x(tokens, 600 + tokens) + core_ref = _clone_core(core) + st = _prefill_state(ctx_lens, n_dec, has_init) + y = _prefill_block(core, x, L, st) + y_ref = _triton_prefill_block(core_ref, x, L, st) + _check(y, y_ref, core, core_ref, f"prefill {ctx_lens}+{n_dec}x1") + + +def test_decode_block_replays_under_cuda_graph() -> None: + """The decode block captured into a CUDA graph after one eager call (the + eager call compiles the decode kernel and sizes its per-batch scratch, as + the engine's warmup does) replays to exactly the eager result.""" + m = 32 + core0, L, x = _core(700), _layer(701), _x(m, 702) + st, _ = _decode_state(m) + eager_core = _clone_core(core0) + y_eager = _target(M.DecodeTarget, eager_core)._gdn_block(x, 0, L, st).clone() + graph_core = _clone_core(core0) + target = _target(M.DecodeTarget, graph_core) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + target._gdn_block(x, 0, L, st) + torch.cuda.current_stream().wait_stream(stream) + torch.cuda.synchronize() + graph_core._bound.conv[0].copy_(core0._bound.conv[0]) + graph_core._bound.ssm[0].copy_(core0._bound.ssm[0]) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + y_graph = target._gdn_block(x, 0, L, st) + assert torch.equal(graph_core._bound.ssm[0], core0._bound.ssm[0]), "capture must not execute" + graph.replay() + torch.cuda.synchronize() + assert torch.equal(y_graph, y_eager) + assert torch.equal(graph_core._bound.ssm[0], eager_core._bound.ssm[0]) + assert torch.equal(graph_core._bound.conv[0], eager_core._bound.conv[0]) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_logits_widening.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_logits_widening.py new file mode 100644 index 000000000000..c369850d3443 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_logits_widening.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The target's logits processor returns the shell's logits bit for bit: the same row gather, +the same fp32 values, its head's bf16 logits widened by one vectorized kernel.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 + modeling as T, +) +from tensorrt_llm._torch.modules.logits_processor import LogitsProcessor + +assert torch.cuda.is_available(), "the logits processor test requires a CUDA device" + +VOCAB, HIDDEN = 248320, 2560 + + +@pytest.fixture(scope="module") +def head(): + g = torch.Generator(device="cuda").manual_seed(0) + lin = torch.nn.Linear(HIDDEN, VOCAB, bias=False, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + lin.weight.copy_(torch.randn(VOCAB, HIDDEN, generator=g, device="cuda") * 0.02) + return lin + + +@pytest.mark.parametrize("seq_lens", [[1] * 128, [5000, 3000, 1] + [1] * 121, [7]]) +def test_matches_the_shell_processor_bitwise(head, seq_lens) -> None: + g = torch.Generator(device="cuda").manual_seed(len(seq_lens)) + lens = torch.tensor(seq_lens, dtype=torch.int32, device="cuda") + hidden = torch.randn(int(lens.sum()), HIDDEN, generator=g, device="cuda").to(torch.bfloat16) + md = SimpleNamespace(seq_lens_cuda=lens) + with torch.no_grad(): + want = LogitsProcessor().forward(hidden, head, md) + got = T._WideningLogitsProcessor().forward(hidden, head, md) + assert got.dtype == torch.float32 and got.shape == (len(seq_lens), VOCAB) + assert torch.equal(got.view(torch.int32), want.view(torch.int32)) + + +def test_context_logits_and_no_metadata(head) -> None: + g = torch.Generator(device="cuda").manual_seed(9) + hidden = torch.randn(33, HIDDEN, generator=g, device="cuda").to(torch.bfloat16) + with torch.no_grad(): + for md, ctx in ((None, False), (SimpleNamespace(seq_lens_cuda=None), True)): + want = LogitsProcessor().forward(hidden, head, md, ctx) + got = T._WideningLogitsProcessor().forward(hidden, head, md, ctx) + assert torch.equal(got.view(torch.int32), want.view(torch.int32)) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_weight_table.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_weight_table.py index 4145ef96632c..60fd360c0352 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_weight_table.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_weight_table.py @@ -65,7 +65,7 @@ ) -def _core(mtp: bool): +def _core(mtp: bool, backend: str = "TRTLLM"): """The configuration surface the table reads.""" return SimpleNamespace( model_config=SimpleNamespace( @@ -75,6 +75,7 @@ def _core(mtp: bool): ), vision_config=SimpleNamespace(**_VISION), mtp_enabled=mtp, + trunk_moe_backend=backend, ) @@ -146,32 +147,34 @@ def hyper_connection(prefix): P(f"{prefix}_up", stream, rank) P(f"{prefix}_inject", t.hc_count, stream) - def gated_attention(prefix): - P(f"{prefix}q_w", 2 * q_width, hidden) - P(f"{prefix}k_w", kv_width, hidden) - P(f"{prefix}v_w", kv_width, hidden) + def gated_attention(prefix, fused_in=False): + if fused_in: + # [q | k | v | gate | index q, k | one index head of zero pad] + rows = 2 * q_width + 2 * kv_width + index_width + t.indexer_head_dim + P(f"{prefix}attn_in_w", rows, hidden) + else: + P(f"{prefix}q_w", 2 * q_width, hidden) + P(f"{prefix}k_w", kv_width, hidden) + P(f"{prefix}v_w", kv_width, hidden) P(f"{prefix}o_w", hidden, q_width) P(f"{prefix}q_norm", t.head_dim) P(f"{prefix}k_norm", t.head_dim) def shared_expert(prefix): - P(f"{prefix}sh_gu", 2 * shared, hidden) + # gate|up, then the per-token gate's row and zero rows (`weights._cat_rows`). + P(f"{prefix}sh_gu", 2 * shared + 8, hidden) P(f"{prefix}sh_dn", hidden, shared) - P(f"{prefix}sh_gate", 1, hidden) for i in range(nl): hyper_connection(f"l{i}_ahc") hyper_connection(f"l{i}_mhc") if t.layer_types[i] == "full_attention": - gated_attention(f"l{i}_") - P(f"l{i}_idx_w", index_width, hidden) + gated_attention(f"l{i}_", fused_in=True) P(f"l{i}_idx_qn", t.indexer_head_dim) P(f"l{i}_idx_kn", t.indexer_head_dim) else: - P(f"l{i}_in_qkv_w", conv_dim, hidden) - P(f"l{i}_in_z_w", value_dim, hidden) - P(f"l{i}_in_a_w", t.linear_num_value_heads, hidden) - P(f"l{i}_in_b_w", t.linear_num_value_heads, hidden) + # [in_proj_qkv | in_proj_z | in_proj_a | in_proj_b], one row block each. + P(f"l{i}_in_proj_w", conv_dim + value_dim + 2 * t.linear_num_value_heads, hidden) P(f"l{i}_A_log", t.linear_num_value_heads) P(f"l{i}_dt_bias", t.linear_num_value_heads) P(f"l{i}_conv_w", conv_dim, 1, t.linear_conv_kernel_dim) @@ -182,6 +185,11 @@ def shared_expert(prefix): P(f"l{i}_fc1_s", experts, 2 * inter, hidden // 16, dtype=u8) P(f"l{i}_fc2_w", experts, hidden, inter // 2, dtype=u8) P(f"l{i}_fc2_s", experts, hidden, inter // 16, dtype=u8) + # The trtllm-gen backend's second, CuteDSL-layout copy of the experts. + P(f"l{i}_fc1_w_cd", experts, 2 * inter, hidden // 2, dtype=u8) + P(f"l{i}_fc1_s_cd", experts, 2 * inter, hidden // 16, dtype=u8) + P(f"l{i}_fc2_w_cd", experts, hidden, inter // 2, dtype=u8) + P(f"l{i}_fc2_s_cd", experts, hidden, inter // 16, dtype=u8) for name in ("e_isc1", "e_isc1_up", "e_ws2_1", "e_ws2_1_up", "e_isc2", "e_ws2_2"): P(f"l{i}_{name}", experts, dtype=f32) shared_expert(f"l{i}_") @@ -260,6 +268,72 @@ def test_the_table_declares_what_the_init_block_did(mtp): assert not wrong, list(wrong.items())[:8] +def test_gdn_input_projections_fill_adjacent_row_blocks(): + """The concatenated GDN input projection: each checkpoint tensor lands in its own + row block, in the column order the forward splits the GEMM output by, and the + blocks tile the declared rows exactly.""" + t = SimpleNamespace(**_TEXT) + core = _core(False) + manifest = W.MODEL_WEIGHTS.manifest(core) + key_dim = t.linear_num_key_heads * t.linear_key_head_dim + value_dim = t.linear_num_value_heads * t.linear_value_head_dim + widths = ( + 2 * key_dim + value_dim, + value_dim, + t.linear_num_value_heads, + t.linear_num_value_heads, + ) + gdn = [i for i, kind in enumerate(t.layer_types) if kind == "linear_attention"] + assert len(gdn) == 36 + for i in gdn: + rows = manifest[f"l{i}_in_proj_w"] + names = [ckpt.rsplit(".", 2)[-2] for ckpt, _, _ in rows] + assert names == ["in_proj_qkv", "in_proj_z", "in_proj_a", "in_proj_b"], names + assert all(ckpt.startswith(f"{W._LM}.layers.{i}.linear_attn.") for ckpt, _, _ in rows) + assert all(transform is None for _, _, transform in rows) + spans = [(index[0].start, index[0].stop) for _, index, _ in rows] + starts = [0, widths[0], widths[0] + widths[1], widths[0] + widths[1] + widths[2]] + assert spans == [(s, s + w) for s, w in zip(starts, widths)], spans + assert spans[-1][1] == sum(widths) == 16480 + + +def test_attention_input_projection_relayout(): + """`_attn_in_w` puts q, k, v, the output gate and the index projection side by side + so that one GEMM's columns are exactly what the four separate projections gave: + q and gate de-interleaved from q_proj's per-head `[q | gate]` rows, then a zero pad + of one index head. Checked at the checkpoint's head geometry, a narrow hidden.""" + t = SimpleNamespace(**_TEXT) + heads, d, kvh, hidden = t.num_attention_heads, t.head_dim, t.num_key_value_heads, 64 + idx_rows = (t.indexer_n_heads + t.indexer_kv_heads) * t.indexer_head_dim + g = torch.Generator().manual_seed(0) + q = torch.randn(2 * heads * d, hidden, generator=g, dtype=torch.float64) + k = torch.randn(kvh * d, hidden, generator=g, dtype=torch.float64) + v = torch.randn(kvh * d, hidden, generator=g, dtype=torch.float64) + idx = torch.randn(idx_rows, hidden, generator=g, dtype=torch.float64) + core = SimpleNamespace(heads_q=heads, head_dim=d, index_head_dim=t.indexer_head_dim) + fused = W._attn_in_w((q, k, v, idx), core) + rows = 2 * heads * d + 2 * kvh * d + idx_rows + t.indexer_head_dim + assert fused.shape == (rows, hidden) and fused.is_contiguous() + + x = torch.randn(3, hidden, generator=g, dtype=torch.float64) + y = x @ fused.t() + qg = (x @ q.t()).reshape(3, heads, 2 * d) + want = torch.cat( + [ + qg[..., :d].reshape(3, -1), + x @ k.t(), + x @ v.t(), + qg[..., d:].reshape(3, -1), + x @ idx.t(), + torch.zeros(3, t.indexer_head_dim, dtype=torch.float64), + ], + 1, + ) + torch.testing.assert_close(y, want, rtol=0, atol=1e-12) + # The row is a whole number of head_dim heads: the main norm+rope addresses it so. + assert rows % d == 0 + + def test_cutedsl_fc1_pairs_up_and_gate_rows_in_64_row_groups(): """Destination rows `[128k, 128k + 64)` hold up rows `[64k, 64k + 64)` and the next 64 hold the matching gate rows, which sit `I` rows later in the `[up ; gate]` stack.""" @@ -285,3 +359,14 @@ def test_cutedsl_fc2_keeps_the_checkpoint_row_order_and_trtllm_gen_shuffles_it() for v in range(4): assert int(shuffled[8 * v + u, 0]) == 4 * u + v assert int(shuffled[32 + 8 * v + u, 0]) == 32 + 4 * u + v + + +@pytest.mark.parametrize("backend", ["TRTLLM", "CUTEDSL"]) +def test_cute_dsl_expert_copy_rows_only_under_trtllm_gen(backend): + """Under trtllm-gen every layer also carries its experts in the CuteDSL layout + (`fc*_cd`), which prefill-sized launches run on; under CuteDSL the one copy is that + layout and no `_cd` row is declared.""" + manifest = W.MODEL_WEIGHTS.manifest(_core(False, backend)) + rows = [k for k in manifest if k.endswith(("_fc1_w_cd", "_fc1_s_cd", "_fc2_w_cd", "_fc2_s_cd"))] + layers = _TEXT["num_hidden_layers"] + assert len(rows) == (4 * layers if backend == "TRTLLM" else 0), len(rows) From 6e34769b5fb97f895deef3ce587a842c8cde0dd9 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Mon, 19 Oct 2026 00:58:37 -0700 Subject: [PATCH 10/28] [TRTLLM-16460][chore] Qwen3.8-Flash-Next target: remove dead paths and duplicated helpers A cleanup pass over the feature and perf commits; nothing the target computes changes. Dead code: - ops/hc: the Triton qwen4_exp_hc_decode_first / _decode_mix / _decode_mix_fp4 entries and qwen4_exp_hc_gate_mean, with their tests. The forward never reached them: one-row decode sites run the CuTe pair, larger batches the fused up projection. The CuTe entries now own the references and checks they borrowed; their NVFP4 reference is catalog fp4_quantize's. - DecodeTarget: the _HC_FUSED_UP_MIN_ROWS / _HC_TC_FIRST_MIN_ROWS branches that _HC_DECODE_GEMV_MAX_ROWS = 1 left unreachable (with the fp4_quantize fallback), the inputs_embeds branch begin_step already refuses, and the unread `pos` / `qsl_long` fields of the decode step tuples. - qsa_sparse_gqa_rowtile: the Triton bitmap and row-tile passes and their launch knobs. The entry now refuses what the CuTe pass cannot take (non-bf16, non-SM103, head_dim != 256, misaligned rows) instead of falling back; the fp16 / fp8 pool cells and the Triton comparison test go. - gdn_prefill_conv_post: the Triton conv pass the CuTe pass always replaced; an unmet precondition is now an assertion. - catalog/ssm/gated_delta_rule_decode_pretranspose and its test: no forward path launched it. Its fp64 recurrence moved to gdn_decode_gated_norm_pdl._recurrence_reference, and the decode tests compare against FlashInfer's kernel directly. - The decoder runner's first-draft mRoPE refusal and its tests: nothing sets py_is_first_draft since #19867. Reuse: - The local TRTLLM_ENABLE_PDL gates call flashinfer_utils.get_env_enable_pdl. - qsa_paged_sparse_gqa_nvfp4 imports the QSA backend's split-K config and PDL policy instead of copying them. - qsa_topk_expand_decode.supports() replaces the width checks restated at its two call sites. - The CuteDSL FC2 entry uses FC1's unswizzle / dequant reference helpers, and gdn_decode_rows_gated_norm's reference delegates to gdn_decode_gated_norm_pdl's. Less work: - gdn_decode_gated_norm_pdl and gdn_in_proj_conv1d_pdl compile from the call's operands instead of dummies spanning the whole coalesced state pool, a transient of every layer's state for every slot during decode graph capture. - PrefillTarget._gdn_core_generic reads q|k|v as the projection's column view and lets the chunked delta rule write `out` directly: two copies fewer per GDN layer on verify and bf16-state steps. Comments no longer cite campaign scratch files. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../modeling_v2/catalog/gemm/low_m_bf16_mm.py | 7 +- .../ssm/flashinfer_chunk_gated_delta_rule.py | 5 +- .../gated_delta_rule_decode_pretranspose.py | 252 --------- .../modeling.py | 141 ++--- .../ops/attention/qsa_paged_index_store.py | 6 +- .../attention/qsa_paged_sparse_gqa_decode.py | 9 - .../attention/qsa_paged_sparse_gqa_nvfp4.py | 37 +- .../attention/qsa_rowtile_attention_cute.py | 23 +- .../ops/attention/qsa_sparse_gqa_rowtile.py | 310 ++--------- .../ops/attention/qsa_topk_expand_decode.py | 20 +- .../qwen4_exp_attn_decode_prologue.py | 8 +- .../ops/attention/qwen4_exp_qk_norm_rope.py | 4 +- .../ops/gemm/bf16_gemv_pdl.py | 15 +- .../ops/hc/qwen4_exp_hc_decode_first.py | 491 ------------------ .../ops/hc/qwen4_exp_hc_decode_first_cute.py | 179 +++++-- .../ops/hc/qwen4_exp_hc_decode_mix.py | 331 ------------ .../ops/hc/qwen4_exp_hc_decode_mix_fp4.py | 431 --------------- .../hc/qwen4_exp_hc_decode_mix_fp4_cute.py | 165 ++++-- .../ops/hc/qwen4_exp_hc_decode_tc_first.py | 15 +- .../ops/hc/qwen4_exp_hc_gate_mean.py | 125 ----- .../ops/hc/qwen4_exp_hc_inject_norm.py | 4 +- .../ops/hc/qwen4_exp_hc_lowrank_act.py | 7 +- .../ops/hc/qwen4_exp_hc_up_gate_mean.py | 12 +- .../ops/hc/qwen4_exp_hc_up_gate_mean_fp4.py | 11 +- .../ops/moe/cute_dsl_nvfp4_grouped_gemm.py | 27 +- .../ops/ssm/gdn_decode_gated_norm_pdl.py | 144 ++--- .../ops/ssm/gdn_decode_rows_gated_norm.py | 85 +-- .../ops/ssm/gdn_in_proj_conv1d_pdl.py | 81 ++- .../ops/ssm/gdn_prefill_conv_post.py | 238 +-------- .../engine/runners/decoder/inputs.py | 13 - .../test_mrope_first_draft_refusal.py | 54 -- .../executor/test_pytorch_model_engine.py | 31 -- ...test_modeling_v2_qsa_sparse_gqa_rowtile.py | 71 +-- ...t_modeling_v2_qwen4_exp_hc_decode_first.py | 450 ---------------- ...eling_v2_qwen4_exp_hc_decode_first_cute.py | 46 +- ...est_modeling_v2_qwen4_exp_hc_decode_mix.py | 157 ------ ...odeling_v2_qwen4_exp_hc_decode_mix_cute.py | 20 +- ...modeling_v2_qwen4_exp_hc_decode_mix_fp4.py | 190 ------- ...ing_v2_qwen4_exp_hc_decode_mix_fp4_cute.py | 42 +- ...test_modeling_v2_qwen4_exp_hc_gate_mean.py | 145 ------ ...t_modeling_v2_qwen4_exp_hc_up_gate_mean.py | 9 +- ...deling_v2_qwen4_exp_hc_up_gate_mean_fp4.py | 20 +- .../test_modeling_v2_cute_dsl_weight_path.py | 13 +- ...t_modeling_v2_ple_gate_conv_inject_norm.py | 5 +- .../ssm/test_modeling_v2_gdn_block.py | 2 +- ...t_modeling_v2_gdn_decode_gated_norm_pdl.py | 20 +- ..._modeling_v2_gdn_decode_rows_gated_norm.py | 20 +- ...v2_gated_delta_rule_decode_pretranspose.py | 167 ------ 48 files changed, 662 insertions(+), 3996 deletions(-) delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/gated_delta_rule_decode_pretranspose.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix_fp4.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_gate_mean.py delete mode 100644 tests/unittest/_torch/executor/test_mrope_first_draft_refusal.py delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_first.py delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix.py delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4.py delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_gate_mean.py delete mode 100644 tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_gated_delta_rule_decode_pretranspose.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/low_m_bf16_mm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/low_m_bf16_mm.py index 72fda36b187e..238af508a3cd 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/low_m_bf16_mm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/low_m_bf16_mm.py @@ -51,11 +51,8 @@ def _s(mma_m: int, mma_n: int, split_k: int, ab_stages: int) -> tuple: #: bound, and any (N, K) not listed, runs on cuBLAS. Every row below is a #: qwen3.8-flash-next-nvfp4 tp1 projection; a target with other shapes adds its #: own rows after measuring them. Each choice beat cublas_mm by >= 5% at the -#: bound's M in -#: qwen3.8-flash-next-nvfp4-perf-opt/05-iterations/candidates/low-m-gemm/microbench/sweep.json -#: or, for the attention in-projection concat and the PLE value projection, -#: .../candidates/lowm-integrate/microbench/sweep_new.json -#: (GB300, cold L2, graph replay, pdl off). M=3 uses the M=4 choice, and so on. +#: bound's M in a sweep of these shapes (GB300, cold L2, graph replay, pdl off). +#: M=3 uses the M=4 choice, and so on. _TABLE: dict[tuple[int, int], tuple[tuple[int, tuple], ...]] = { # GDN in_qkv: 16 q + 16 k heads x 128, 48 v heads x 128 (10240, 2560): ((1, _d(64, 1)), (2, _d(64, 2)), (8, _s(128, 8, 1, 6))), diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/flashinfer_chunk_gated_delta_rule.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/flashinfer_chunk_gated_delta_rule.py index 7504fb20757a..f4d3e1e69849 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/flashinfer_chunk_gated_delta_rule.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/flashinfer_chunk_gated_delta_rule.py @@ -147,8 +147,9 @@ class FlashinferChunkGatedDeltaRule(OpWrapper): contiguous `[1, T, H, 128]`; `g` fp32 log-space; `beta` fp32). Give it context rows only: single-token rows are slow here. Measured per layer at 8065 + 127 x 1 tokens: 1078.8 us with every row in this op, 399.4 us - with the 127 decode rows sent to `gated_delta_rule_decode_pretranspose` - instead (both halves plus the concatenation). + with the 127 decode rows sent to FlashInfer's decode kernel + (`gdn_decode.gated_delta_rule_decode_pretranspose`) instead (both halves + plus the concatenation). Memory per call: a packed fp32 copy of the named slots in and out (2 x 3 MiB per sequence), `exp(g)`, and the output. q/k/v/beta are not copied diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/gated_delta_rule_decode_pretranspose.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/gated_delta_rule_decode_pretranspose.py deleted file mode 100644 index 20a11f7cf5db..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/gated_delta_rule_decode_pretranspose.py +++ /dev/null @@ -1,252 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Gated DeltaNet single-token decode over an fp32 slot pool, gates fused in. - -Wraps FlashInfer's public -``flashinfer.gdn_decode.gated_delta_rule_decode_pretranspose`` (tier-3 -sourcing, the installed FlashInfer). With an fp32 pool and ``T == 1`` it -runs the CuTe-DSL "pretranspose" kernel (``run_pretranspose_decode``): eight -CTAs per (sequence, v-head) state, each streaming 16 of its 128 V rows through -shared memory, the pool read and written in place through -``initial_state_indices``. A bf16 pool or ``T > 1`` route elsewhere inside the -same callable and are not certified here. - -FlashInfer API: ``flashinfer.gdn_decode.gated_delta_rule_decode_pretranspose``. -""" - -from __future__ import annotations - -from typing import Optional - -import flashinfer.gdn_decode -import torch - -from .._op import Arch, Cell, OpWrapper - -#: qwen3.8-flash-next-nvfp4's Gated DeltaNet at tp1: (H, HV, K, V) -- 16 q/k -#: heads, 48 v heads, head dims 128. -QWEN38_GEOMETRY = (16, 48, 128, 128) -#: Largest slot count the cache manager was seen to allocate the recurrent pool -#: at for qwen3.8-flash-next-nvfp4 (10 .. 461 over twelve runtime-value dumps). -ENGINE_SLOTS = 461 - - -def _scattered(n: int, slots: int = ENGINE_SLOTS) -> tuple[int, ...]: - """`n` distinct slots of a `slots`-row pool, both ends first, then a - stride-37 walk over the interior (37 is coprime to every pool size here).""" - rest = tuple((i * 37 + 4) % (slots - 2) + 1 for i in range(max(n - 2, 0))) - return ((slots - 1, 0) + rest)[:n] - - -class GatedDeltaRuleDecodePretranspose(OpWrapper): - """One gated-delta-rule step per sequence, gates computed in-kernel. - - Row `n` of `q`/`k` `[B, 1, H, K]`, `v` `[B, 1, HV, V]`, `a`/`b` `[B, 1, HV]`; - v-head `hv` reads q/k head `hv // (HV // H)`; state `h` `[V, K]` (K - innermost) read from `initial_state[initial_state_indices[n]]`, in fp32: - - g = -exp(A_log) * softplus(a + dt_bias) (softplus beta 1, threshold 20) - beta = sigmoid(b) - q, k = x * rsqrt(sum(x^2) + 1e-6) (only with use_qk_l2norm) - h *= exp(g); h += k (beta * (v - h k))^T; o = h (scale * q) - - `scale=None` is `K ** -0.5`. The final `h` goes back to the same slot, or - to `output_state_indices[n]` when given. A negative index skips the row: - no state read or write, `o` written as zero. - - Fusion boundary: both gates, the optional q/k norm, scale, recurrence, - indexed gather and scatter. The conv, the q/k/v split and the gated RMS - norm on `o` are the caller's. - - Returns `(output, initial_state)` as the callable does: `output` - `[B, 1, HV, V]` (the `output` buffer itself when given, bf16; else a fresh - tensor in `q.dtype`) and the pool, advanced in place. - """ - - ARCHS = frozenset({Arch.SM_103}) - - # spec: geometry (H, HV, K, V), m (rows), slots (pool rows), indices (None - # -> _scattered(m, slots)), qkv ("views": column views of one - # [m, 2*H*K + HV*V] conv output, as a decode forward splits it | - # "contiguous"), output ("buffer" | None), pool ("contiguous" | "padded": - # slot stride HV*V*K + 64 | "layer_view": pool[1] of a [3, slots, ...] - # cache). Every cell but the last is qwen3.8-flash-next-nvfp4's geometry. - CELLS: tuple[Cell, ...] = tuple( - Cell( - why=f"qwen3.8 decode M={m}, 461-slot pool, scattered slots", - spec=dict(geometry=QWEN38_GEOMETRY, m=m, slots=ENGINE_SLOTS), - ) - for m in (1, 2, 8, 32, 64, 128) - ) + ( - Cell( - why="qwen3.8 output None: fresh q.dtype output, M=4", - spec=dict(geometry=QWEN38_GEOMETRY, m=4, slots=ENGINE_SLOTS, output=None), - ), - Cell( - why="qwen3.8 CUDA-graph pad rows (-1) between slot 0 and 460", - spec=dict( - geometry=QWEN38_GEOMETRY, - m=4, - slots=ENGINE_SLOTS, - indices=(0, -1, ENGINE_SLOTS - 1, -1), - ), - ), - Cell( - why="qwen3.8 padded slot stride (HV*V*K + 64), M=8", - spec=dict(geometry=QWEN38_GEOMETRY, m=8, slots=40, pool="padded"), - ), - Cell( - why="qwen3.8 pool[1] of a [3, slots, HV, V, K] cache, M=4", - spec=dict(geometry=QWEN38_GEOMETRY, m=4, slots=24, pool="layer_view"), - ), - Cell( - why="qwen3.8 contiguous q/k/v instead of conv-output views, M=16", - spec=dict(geometry=QWEN38_GEOMETRY, m=16, slots=ENGINE_SLOTS, qkv="contiguous"), - ), - Cell( - why="HV/H = 2 (16 q/k heads, 32 v heads), M=8, 40-slot pool", - spec=dict(geometry=(16, 32, 128, 128), m=8, slots=40), - ), - ) - - note = """ - Metadata consumed: none. - - Call site as qwen3.8-flash-next-nvfp4's decode `_gdn_block` makes it (in - place of `fused_sigmoid_gating_delta_rule_update`): `q=reshape(qv, [M, 1, 16, 128])`, - `k=...`, `v=reshape(vv, [M, 1, 48, 128])` -- column views of the conv - output, read through their strides -- `state=None`, `A_log`/`dt_bias` - fp32 `[48]`, `a`/`b` the `[M, 48]` projections viewed `[M, 1, 48]`, - `scale=None`, `output=` a bf16 `[M, 1, 48, 128]` buffer, `use_qk_l2norm= - True`, `initial_state=` the layer's fp32 pool `[slots, 48, 128, 128]`, - `initial_state_indices=` int32 slots. `output.view(1, M, 48, 128)` is what - `_gdn_epilogue` reads. - - Preconditions not checked by anything: - - fp32 pool, K-innermost `(V, K)` trailing axes, slot stride a multiple of - 4 elements (asserted) -- a transposed `[S, HV, K, V]` pool at K == V has - the same shape and is accepted. - - Indices distinct and in `[0, S)` or negative. An index >= S is an - unchecked device access: measured, slot 460 of a 40-slot pool raised - `cudaErrorIllegalAddress` on the next sync and poisoned the process. - - `A_log` fp32 (asserted), `dt_bias` fp32 or bf16 (asserted). - - Loud: every tensor's base pointer must be 16-byte aligned (the kernel is - compiled with `assumed_align=16`). Measured: indices taken as - `idx[n_c:]` of an int32 vector at an offset not a multiple of 4 raise - `RuntimeError: Tensor data pointer is not aligned to 16 bytes.` -- so a - mixed step's decode half, `si_all` split after its context rows, needs - a copy. The decode-only step's `si_all` starts at offset 0. - - JIT and capture: the first call per (`T, H, HV, K, V, q.dtype, scale, - use_qk_l2norm`, pool inner strides) compiles a CuTe-DSL kernel in-process - (seconds); the first call per batch size allocates two small int32 aux - tensors, unused on the pool path. Both must happen in eager warmup, before - graph capture. Capture-safe after that: no host sync, no allocation when - `output` is given (one `torch.zeros` per call otherwise). - - Numerics (measured at qwen3.8's geometry, fp64 reference): o max abs 1.2e-4 - at M=128 (|o| <= 6.2e-2), equal to the in-tree Triton op's - (`modules/fla/fused_sigmoid_gating_recurrent`); written state max abs - 2.6e-7 (Triton 2.0e-7). Not bitwise to the Triton op: rsqrt(sum + eps) vs - 1 / (sqrt(sum) + eps) for the q/k norm, and fastmath exp/log. - - Not certified: bf16/fp16 pool (another kernel), `T > 1` (routes to - `gated_delta_rule_mtp`), `state=` direct path, `output_state_indices`, - fp16 operands, K or V other than 128. - - Hard limits - none probed. - """ - - def __call__( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - state: Optional[torch.Tensor], - A_log: torch.Tensor, - a: torch.Tensor, - dt_bias: torch.Tensor, - b: torch.Tensor, - scale: Optional[float] = None, - output: Optional[torch.Tensor] = None, - use_qk_l2norm: bool = True, - initial_state: Optional[torch.Tensor] = None, - initial_state_indices: Optional[torch.Tensor] = None, - output_state_indices: Optional[torch.Tensor] = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - return flashinfer.gdn_decode.gated_delta_rule_decode_pretranspose( - q=q, - k=k, - v=v, - state=state, - A_log=A_log, - a=a, - dt_bias=dt_bias, - b=b, - scale=scale, - output=output, - use_qk_l2norm=use_qk_l2norm, - initial_state=initial_state, - initial_state_indices=initial_state_indices, - output_state_indices=output_state_indices, - ) - - def reference( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - state: Optional[torch.Tensor], - A_log: torch.Tensor, - a: torch.Tensor, - dt_bias: torch.Tensor, - b: torch.Tensor, - scale: Optional[float] = None, - output: Optional[torch.Tensor] = None, - use_qk_l2norm: bool = True, - initial_state: Optional[torch.Tensor] = None, - initial_state_indices: Optional[torch.Tensor] = None, - output_state_indices: Optional[torch.Tensor] = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - """`(o, pool_after)`: fp64, rounded once -- `o` to `output`'s dtype (bf16 - when given, else `q.dtype`), the state to the pool's. `state=` (no pool) - is treated as a pool of `B` rows read and written in order. Mutates - nothing.""" - dt = torch.float64 - batch, _, num_k_heads, k_dim = q.shape - hv = v.shape[2] - rep = hv // num_k_heads - scale = k_dim**-0.5 if scale is None else scale - pool_in = initial_state if initial_state is not None else state - if initial_state is not None: - reads = initial_state_indices.tolist() - else: - reads = list(range(batch)) - writes = reads if output_state_indices is None else output_state_indices.tolist() - pool = pool_in.clone() - o = torch.zeros(batch, 1, hv, v.shape[3], dtype=dt, device=q.device) - neg_a = -torch.exp(A_log.to(dt)) - for n in range(batch): - if reads[n] < 0: - continue - q_n = q[n, 0].to(dt).repeat_interleave(rep, dim=0) - k_n = k[n, 0].to(dt).repeat_interleave(rep, dim=0) - if use_qk_l2norm: - q_n = q_n * torch.rsqrt(q_n.pow(2).sum(-1, keepdim=True) + 1e-6) - k_n = k_n * torch.rsqrt(k_n.pow(2).sum(-1, keepdim=True) + 1e-6) - x = a[n, 0].to(dt) + dt_bias.to(dt) - softplus = torch.where(x <= 20.0, torch.log1p(torch.exp(x)), x) - beta = torch.sigmoid(b[n, 0].to(dt)) - h = pool_in[reads[n]].to(dt) * torch.exp(neg_a * softplus)[:, None, None] - u = (v[n, 0].to(dt) - torch.einsum("hvk,hk->hv", h, k_n)) * beta[:, None] - h = h + u[:, :, None] * k_n[:, None, :] - o[n, 0] = torch.einsum("hvk,hk->hv", h, q_n * scale) - if writes[n] >= 0: - pool[writes[n]] = h.to(pool.dtype) - out_dtype = output.dtype if output is not None else q.dtype - return o.to(out_dtype), pool - - -gated_delta_rule_decode_pretranspose = GatedDeltaRuleDecodePretranspose() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py index 6496ccd141bb..4f7b5f2e4a3c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py @@ -27,9 +27,9 @@ weight, ``ops/hc/qwen4_exp_hc_lowrank_act`` and ``ops/hc/qwen4_exp_hc_up_gate_mean`` (the up projection as a tcgen05 GEMM whose epilogue applies the sigmoid gate and the stream mean). ``DecodeTarget`` mixers run the first projection as one tcgen05 launch - (``ops/hc/qwen4_exp_hc_decode_tc_first``, ``_HC_TC_FIRST_MIN_ROWS`` to - ``_HC_TC_FIRST_MAX_ROWS`` rows) or, at ``_HC_DECODE_GEMV_MAX_ROWS`` rows, two GEMV - launches (``ops/hc/qwen4_exp_hc_decode_first_cute``, ``ops/hc/qwen4_exp_hc_decode_mix``); + (``ops/hc/qwen4_exp_hc_decode_tc_first``, up to ``_HC_TC_FIRST_MAX_ROWS`` rows) or, at + ``_HC_DECODE_GEMV_MAX_ROWS`` rows, the whole site as two GEMV launches + (``ops/hc/qwen4_exp_hc_decode_first_cute``, ``ops/hc/qwen4_exp_hc_decode_mix_fp4_cute``); at the 48 MoE sites the mixer's last launch also writes the routed experts' NVFP4 input. The MTP layer and the prefill PLE block keep the unfused ``qwen4_exp_gated_residual`` / ``qwen4_exp_post_inject`` / ``_grouped_norm``. @@ -128,9 +128,6 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog.graph.eager_on_graph import ( eager_on_graph, ) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( - fp4_quantize, -) from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.flashinfer_chunk_gated_delta_rule import ( # noqa: E501 flashinfer_chunk_gated_delta_rule, ) @@ -165,6 +162,7 @@ from .ops.attention.qsa_prefill_compress import qsa_prefill_compress from .ops.attention.qsa_sparse_gqa_rowtile import qsa_sparse_gqa_rowtile from .ops.attention.qsa_topk_expand_decode import qsa_topk_expand_decode +from .ops.attention.qsa_topk_expand_decode import supports as _topk_expand_supports from .ops.attention.qsa_topk_select import qsa_topk_select from .ops.attention.qwen4_exp_attn_decode_prologue import qwen4_exp_attn_decode_prologue from .ops.attention.qwen4_exp_qk_norm_rope import qwen4_exp_qk_norm_rope @@ -174,13 +172,11 @@ from .ops.gemm.bf16_gemv_pdl import bf16_gemv_pdl from .ops.hc.qwen4_exp_gated_residual import qwen4_exp_gated_residual from .ops.hc.qwen4_exp_hc_decode_first_cute import qwen4_exp_hc_decode_first_cute -from .ops.hc.qwen4_exp_hc_decode_mix import qwen4_exp_hc_decode_mix from .ops.hc.qwen4_exp_hc_decode_mix_fp4_cute import ( qwen4_exp_hc_decode_mix_cute, qwen4_exp_hc_decode_mix_fp4_cute, ) from .ops.hc.qwen4_exp_hc_decode_tc_first import qwen4_exp_hc_decode_tc_first -from .ops.hc.qwen4_exp_hc_gate_mean import qwen4_exp_hc_gate_mean from .ops.hc.qwen4_exp_hc_inject_norm import qwen4_exp_hc_inject_norm from .ops.hc.qwen4_exp_hc_lowrank_act import qwen4_exp_hc_lowrank_act from .ops.hc.qwen4_exp_hc_up_gate_mean import qwen4_exp_hc_up_gate_mean @@ -454,10 +450,6 @@ def _decode_projection(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor # 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 @@ -544,26 +536,17 @@ def _ple_host_offload_enabled() -> bool: _MTP_EMBED_BLOCK_ALL_STREAMS = True #: Decode batches up to this many rows run each trunk hyper-connection site as -#: `qwen4_exp_hc_decode_first_cute` + `qwen4_exp_hc_decode_mix` (two GEMV launches); -#: larger ones as `qwen4_exp_hc_inject_norm` + cuBLAS. A dispatch threshold from the -#: GB300 micro-bench (the GEMVs win at 1 row and lose from 2), at most the entries' -#: MAX_ROWS. +#: `qwen4_exp_hc_decode_first_cute` + `qwen4_exp_hc_decode_mix_cute` (or its NVFP4 +#: variant at an MoE site; two GEMV launches); larger ones as `qwen4_exp_hc_inject_norm` +#: and the mixer. A dispatch threshold from the GB300 micro-bench (the GEMVs win at 1 +#: row and lose from 2), at most the entries' MAX_ROWS. _HC_DECODE_GEMV_MAX_ROWS = 1 # capacity: fixed: dispatch threshold, not a buffer bound -#: Decode row count from which the hyper-connection up projection runs as -#: `qwen4_exp_hc_up_gate_mean` (gate and stream mean in the GEMM epilogue). The -#: fused kernel's 80 tiles read the 6.5 MB up weight slower than cuBLAS at -#: one row (5.8 vs 5.2 us); from 2 rows on the saved launch and logits round -#: trip win (5.6 vs 6.1 us at 2, 5.5 vs 7.4 at 128; GB300, graph-replayed, -#: cold weights). Prefill always takes the fused op. Trunk sites at most -#: `_HC_DECODE_GEMV_MAX_ROWS` rows never reach the mixer (two GEMV launches). -_HC_FUSED_UP_MIN_ROWS = 2 # capacity: fixed: dispatch threshold, not a buffer bound - -#: Decode row counts whose hyper-connection mixer runs its first half (projection -#: and activations) as the one tcgen05 launch `qwen4_exp_hc_decode_tc_first` -#: instead of cuBLAS split-K (+ its reduce) + `qwen4_exp_hc_lowrank_act`. The upper -#: bound is the entry's MAX_ROWS (one 128-row MMA tile); larger batches keep cuBLAS. -_HC_TC_FIRST_MIN_ROWS = 2 # capacity: fixed: dispatch threshold, not a buffer bound +#: Decode row counts up to which a hyper-connection mixer (so from +#: `_HC_DECODE_GEMV_MAX_ROWS + 1` rows) runs its first half (projection and +#: activations) as the one tcgen05 launch `qwen4_exp_hc_decode_tc_first` instead of +#: cuBLAS split-K (+ its reduce) + `qwen4_exp_hc_lowrank_act`: the entry's MAX_ROWS +#: (one 128-row MMA tile); larger batches keep cuBLAS. _HC_TC_FIRST_MAX_ROWS = ( 128 # capacity: fixed: dispatch threshold (one MMA tile), not a buffer bound ) @@ -2991,14 +2974,13 @@ def _attention_core(self, proj, proj_qkv, index_scratch, attn_out, i) -> None: if num_ctx < num_tokens: parts.append((num_ctx, num_tokens, False)) # Each part's canonical block top-k and its expansion are one launch writing its - # rows of the step's selection (QSA's row starts are zero); a score table wider - # than that kernel's 4096 columns takes the composition. + # rows of the step's selection (QSA's row starts are zero); a score table that + # kernel does not serve (wider than its 4096 columns) takes the composition. selected = torch.empty( [num_tokens, core.token_topk + core.compress_ratio - 1], dtype=torch.int32, device=proj.device, ) - fused_select = block_table.shape[1] * tokens_per_block // core.compress_ratio <= 4096 for lo, hi, context_rows in parts: scores = qsa_paged_index_scores( q=q_index[lo:hi], @@ -3011,7 +2993,7 @@ def _attention_core(self, proj, proj_qkv, index_scratch, attn_out, i) -> None: only_visible_blocks=True, context_rows=context_rows, ) - if fused_select and scores.shape[1] % 4 == 0: + if _topk_expand_supports(scores): qsa_topk_expand_decode( scores, visible_blocks[lo:hi], @@ -3308,9 +3290,10 @@ def _gdn_core(self, proj, out, i) -> None: def _gdn_core_generic(self, proj, out, i) -> None: """`_gdn_core` on a verify step or over a bf16 recurrent pool: the conv, the post-conv gating and the chunked delta rule over the step's `n` rows, or under - speculation the verify path for the generation rows. The q|k|v, a and b columns are - copied out of the projection first: these entries are certified on contiguous - operands.""" + speculation the verify path for the generation rows. The q|k|v columns are read in + place, as the column view of the projection every consumer here is certified on; + the a and b columns are copied out first, as the verify path reads them contiguous. + The chunked delta rule writes its rows of `out` directly.""" core = self.core L = self._layers[i] step = self._step @@ -3334,7 +3317,7 @@ def _gdn_core_generic(self, proj, out, i) -> None: ) = step.gdn num_tokens = step.num_tokens mixed_qkv, _, a, b = torch.split(proj[:num_tokens], core.in_proj_split, 1) - mixed_qkv, a, b = mixed_qkv.contiguous(), a.contiguous(), b.contiguous() + a, b = a.contiguous(), b.contiguous() conv_states = core._bound.conv[i] ssm_states = core._bound.ssm[i] @@ -3379,7 +3362,7 @@ def _gdn_core_generic(self, proj, out, i) -> None: softplus_threshold=_SOFTPLUS_THRESHOLD, ) # q and k are already L2-normalized by the post-conv call. - out_p, _ = chunk_gated_delta_rule( + chunk_gated_delta_rule( q, k, v, @@ -3392,12 +3375,11 @@ def _gdn_core_generic(self, proj, out, i) -> None: output_final_state=False, cu_seqlens=qsl_long_p, use_qk_l2norm_in_kernel=False, - output=None, + output=out[:, :num_prefill_tokens], ) out_d = self._gdn_verify( mq_d, a_d, b_d, ssm_states, i, L, si_d, int_idx, num_decodes, gen_p, replay ) - out[:, :num_prefill_tokens].copy_(out_p) out[:, num_prefill_tokens:num_tokens].copy_(out_d) out[:, num_tokens:].zero_() return @@ -3454,7 +3436,7 @@ def _gdn_core_generic(self, proj, out, i) -> None: softplus_threshold=_SOFTPLUS_THRESHOLD, ) # q and k are already L2-normalized; the recurrence must not repeat it. - recurrence, _ = chunk_gated_delta_rule( + chunk_gated_delta_rule( q, k, v, @@ -3467,9 +3449,8 @@ def _gdn_core_generic(self, proj, out, i) -> None: output_final_state=False, cu_seqlens=qsl_long, use_qk_l2norm_in_kernel=False, - output=None, + output=out[:, :num_tokens], ) - out[:, :num_tokens].copy_(recurrence) out[:, num_tokens:].zero_() def _gdn_prefill_prep(self, mq, a, b, L, i, qsl, si, hi): @@ -4423,12 +4404,12 @@ def step_args(self, md: TrtllmAttentionMetadata) -> dict: def _hc_first(self, normed, first, with_inject): """The mixer's first half from the site's normed stream: the projection over `[down; block_inject; 0]` and the SiLU / injection activations. - Returns `(low, inject)`. Between `_HC_TC_FIRST_MIN_ROWS` and - `_HC_TC_FIRST_MAX_ROWS` rows (a per-capture constant) one tcgen05 launch, - `qwen4_exp_hc_decode_tc_first`; otherwise cuBLAS's fp32 logits and - `qwen4_exp_hc_lowrank_act`.""" + Returns `(low, inject)`. Up to `_HC_TC_FIRST_MAX_ROWS` rows (a per-capture + constant) one tcgen05 launch, `qwen4_exp_hc_decode_tc_first`; larger batches + cuBLAS's fp32 logits and `qwen4_exp_hc_lowrank_act`. Only sites past + `_HC_DECODE_GEMV_MAX_ROWS` rows reach a mixer.""" core = self.core - if _HC_TC_FIRST_MIN_ROWS <= normed.shape[0] <= _HC_TC_FIRST_MAX_ROWS: + if normed.shape[0] <= _HC_TC_FIRST_MAX_ROWS: return qwen4_exp_hc_decode_tc_first( normed, first, core.hc_lowrank, core.hc_count, with_inject ) @@ -4443,21 +4424,16 @@ def _hc_first(self, normed, first, with_inject): def _hc_mixer(self, normed, first, up, with_inject): """One hyper-connection mixer from the site's normed stream: the first projection over `[down; block_inject; 0]` (fp32 logits), the SiLU and - injection activations, the up projection, and the sigmoid gate with the - stream mean. From `_HC_FUSED_UP_MIN_ROWS` rows (a per-capture constant) - the up projection carries the gate and the mean in its epilogue; below - it cuBLAS's up projection + `qwen4_exp_hc_gate_mean` is faster. - Returns `(mixed, inject)`; `inject` is None without `with_inject` (the - stack-final mixer). + injection activations, then the up projection with the sigmoid gate and the + stream mean in its epilogue. Returns `(mixed, inject)`; `inject` is None + without `with_inject` (the stack-final mixer). Neither projection takes the step's `gemm` (`low_m_bf16_mm` under CUDA graphs): neither FlashInfer low-M kernel writes the first projection's fp32 logits, and the up projection's K = 320 tiles neither of them.""" core = self.core low, inject = self._hc_first(normed, first, with_inject) - if normed.shape[0] >= _HC_FUSED_UP_MIN_ROWS: - return qwen4_exp_hc_up_gate_mean(low, up, normed, core.hc_count), inject - return qwen4_exp_hc_gate_mean(cublas_mm(low, up), normed, core.hc_count), inject + return qwen4_exp_hc_up_gate_mean(low, up, normed, core.hc_count), inject def _hc_site(self, hc, block_out, inject, norm, first, up, with_inject, normed=None, moe=None): """One trunk hyper-connection site: the previous block's injection (none at @@ -4491,8 +4467,9 @@ def _hc_site(self, hc, block_out, inject, norm, first, up, with_inject, normed=N topk_w, topk_rows, ) - mix = qwen4_exp_hc_decode_mix_cute if hc.shape[0] == 1 else qwen4_exp_hc_decode_mix - mixed, inject = mix(part, ss, up, hc, norm, core.eps, core.hc_count, with_inject) + mixed, inject = qwen4_exp_hc_decode_mix_cute( + part, ss, up, hc, norm, core.eps, core.hc_count, with_inject + ) return hc, mixed, inject if normed is None: hc, normed = qwen4_exp_hc_inject_norm( @@ -4514,10 +4491,9 @@ def _hc_moe_site(self, hc, block_out, inject, norm, first, up, gscale): """`_hc_site` at an MoE site (always with injection), returning the routed experts' NVFP4 input (linear scales, the layer's global scale `gscale`) beside the mix, written by the mixer's own last launch: on the - GEMV route `qwen4_exp_hc_decode_mix_fp4_cute`, from `_HC_FUSED_UP_MIN_ROWS` - rows the fused up projection `qwen4_exp_hc_up_gate_mean_fp4`. Between - the two (only if the thresholds leave a gap) `fp4_quantize` reads the - mix back. Returns `(hc, mixed, inject, xq, xsf)`.""" + GEMV route `qwen4_exp_hc_decode_mix_fp4_cute`, otherwise the fused up + projection `qwen4_exp_hc_up_gate_mean_fp4`. Returns + `(hc, mixed, inject, xq, xsf)`.""" core = self.core if hc.shape[0] <= _HC_DECODE_GEMV_MAX_ROWS: hc, part, ss = qwen4_exp_hc_decode_first_cute( @@ -4528,12 +4504,8 @@ def _hc_moe_site(self, hc, block_out, inject, norm, first, up, gscale): ) return hc, mixed, inject, xq, xsf hc, normed = qwen4_exp_hc_inject_norm(hc, block_out, inject, norm, core.eps, core.hc_count) - if normed.shape[0] >= _HC_FUSED_UP_MIN_ROWS: - low, inject = self._hc_first(normed, first, True) - mixed, xq, xsf = qwen4_exp_hc_up_gate_mean_fp4(low, up, normed, core.hc_count, gscale) - return hc, mixed, inject, xq, xsf - mixed, inject = self._hc_mixer(normed, first, up, True) - xq, xsf = fp4_quantize(mixed, gscale, _SF_VEC, False, _SF_LINEAR) + low, inject = self._hc_first(normed, first, True) + mixed, xq, xsf = qwen4_exp_hc_up_gate_mean_fp4(low, up, normed, core.hc_count, gscale) return hc, mixed, inject, xq, xsf def _grouped_norm(self, x: torch.Tensor, weights: tuple) -> torch.Tensor: @@ -4616,7 +4588,7 @@ def _attention_prologue( attention core joins it before reading either.""" core = self.core num_tokens = proj.shape[0] - req_idx, logical, _, _, _, _, block_table, tokens_per_block, _, pos3, coords, _ = qsa + req_idx, logical, _, _, _, _, block_table, tokens_per_block, pos3, coords, _ = qsa side.wait_stream(torch.cuda.current_stream()) qwen4_exp_qk_norm_rope( proj, @@ -4718,7 +4690,6 @@ def _attention_block( topk_row_starts, block_table, tokens_per_block, - pos, pos3, coords, sparse, @@ -4817,7 +4788,7 @@ def _attention_block( use_pdl=fused_prologue, prefetch=fused_prologue and num_tokens <= _SPARSE_DECODE_MAX_ROWS, ) - if scores.shape[1] % 4 == 0 and scores.shape[1] <= 4096: + if _topk_expand_supports(scores): # The canonical block top-k and its expansion to tokens in one launch # (QSA's row starts are zero). selected = qsa_topk_expand_decode( @@ -4988,7 +4959,7 @@ def _gdn_block(self, x, i, L, gdn_state, gemm=cublas_mm, fused_conv=False) -> to Over a bf16 pool the recurrence is `fused_sigmoid_gating_delta_rule_update`'s rows form, which is batch-invariant.""" core = self.core - si_all, qsl_long, is_verify, num_decodes, gen_p, int_idx, replay = gdn_state + si_all, is_verify, num_decodes, gen_p, int_idx, replay = gdn_state num_tokens = x.shape[0] ssm_states = core._bound.ssm[i] fp32_state = ssm_states.dtype == torch.float32 @@ -5221,7 +5192,6 @@ def forward( attn_metadata: AttentionMetadata, input_ids: torch.IntTensor | None = None, position_ids: torch.IntTensor | None = None, - inputs_embeds: torch.FloatTensor | None = None, **kwargs, ): core = self.core @@ -5234,9 +5204,8 @@ def forward( # (`_decode_projection`); an eager step keeps cuBLAS. gemm = _decode_projection if md.is_cuda_graph else cublas_mm - # Three views of one `[3, 1, T]` buffer (see PrefillTarget.forward): `pos`, - # contiguous `pos3`, and `coords`. - pos = torch.reshape(torch.split(position_ids, [1, _QSA_POSITION_AXES - 1], 0)[0], [-1]) + # Two views of one `[3, 1, T]` buffer (see PrefillTarget.forward): contiguous + # `pos3` and `coords`. coords = torch.transpose(torch.reshape(position_ids, [_QSA_POSITION_AXES, -1]), 0, 1) pos3 = torch.empty( [_QSA_POSITION_AXES, position_ids.shape[-1]], @@ -5245,18 +5214,7 @@ def forward( ) pos3.copy_(torch.reshape(position_ids, [_QSA_POSITION_AXES, -1])) - # Spliced vision rows are already the embedding; a lookup would overwrite them. - if inputs_embeds is None: - h = nn.functional.embedding(input_ids, core.w["embed"]) - else: - assert inputs_embeds.shape[0] == input_ids.shape[0], ( - f"the spliced embedding carries {inputs_embeds.shape[0]} rows for " - f"{input_ids.shape[0]} token ids" - ) - assert inputs_embeds.shape[1] == core.hidden, ( - f"the spliced embedding is {inputs_embeds.shape[1]} wide, not {core.hidden}" - ) - h = inputs_embeds + h = nn.functional.embedding(input_ids, core.w["embed"]) num_tokens = h.shape[0] dt = h.dtype dev = h.device @@ -5311,7 +5269,7 @@ def forward( int_idx = core._bound.int_idx[:batch] if core._bound.replay is not None: replay = _replay_work(core, mm, batch) - gdn_state = (si_all, qsl_long, is_verify, batch, gen_p, int_idx, replay) + gdn_state = (si_all, is_verify, batch, gen_p, int_idx, replay) # Every generation sequence carries its n-gram history; the engine's # `has_initial_states` covers the context prefix only (empty here). @@ -5346,7 +5304,6 @@ def forward( md.qsa_topk_row_starts[:num_tokens], md.qsa_block_table, md.tokens_per_block, - pos, pos3, coords, sparse, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_store.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_store.py index 918cb5f91716..a201309b62df 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_store.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_store.py @@ -109,7 +109,7 @@ def _qsa_paged_index_store_kernel( def _launch_rows(rows: int) -> Tuple[int, int]: """Rows per program and warps: one row per program while the grid is small, a tile of rows once there are enough rows to fill the machine. Measured on - GB300 against 1/4/8/16/32 rows per program (microbench `qsa-index-store-kernel`).""" + GB300 against 1/4/8/16/32 rows per program.""" if rows <= 256: return 1, 1 return 16, 4 @@ -140,7 +140,7 @@ def _launch_rows(rows: int) -> Tuple[int, int]: _DECODE = ((0, 1), (1, 1), (31, 1), (32, 1), (33, 1), (63, 1), (64, 1), (95, 1)) -#: The campaign's envelope: ISL 8192 + OSL 1024 puts decode rows up to position +#: The target's serving envelope: ISL 8192 + OSL 1024 puts decode rows up to position #: 9215 (page column 287), 128 concurrent requests, and `max_num_tokens` 16384 #: rows in one prefill or mixed step. _WIDE = dict(slots=160, width=300) @@ -279,7 +279,7 @@ class QsaPagedIndexStore(OpWrapper): why="int64 requests and int64 positions", spec=_spec(_DECODE, pages=32, request_dtype=torch.int64), ), - # ---- the campaign envelope: decode M=128, 16384-row steps ---------- + # ---- the serving envelope: decode M=128, 16384-row steps ----------- Cell( why="decode M=128 at positions up to 9153 (page column 286)", spec=_spec(_DECODE_128, pages=18432, **_WIDE), diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_decode.py index 33274e1c2b2b..d53be3643322 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_decode.py @@ -39,7 +39,6 @@ import functools import math -import os from typing import Optional import cuda.bindings.driver as cuda_drv @@ -82,13 +81,6 @@ LOG2E = 1.4426950408889634 -@functools.lru_cache(maxsize=None) -def pdl_enabled(device_index: int) -> bool: - """Programmatic dependent launch: sm_90+ unless TRTLLM_ENABLE_PDL=0.""" - props = torch.cuda.get_device_properties(device_index) - return props.major >= 9 and os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" - - class _SplitKKernel: def __init__( self, @@ -623,7 +615,6 @@ def launch( if rows == 0: return output dev = q.device - pdl = pdl and pdl_enabled(dev.index if dev.index is not None else torch.cuda.current_device()) use_pdl = bool(pdl and _qsa_pdl_enabled(rows)) partial_output = torch.empty( (n_splits, rows, num_q_heads, head_dim), dtype=torch.float32, device=dev diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4.py index 23f4f7bb3fee..95fa751a5f84 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4.py @@ -4,7 +4,6 @@ from __future__ import annotations -import os from typing import Optional import torch @@ -12,7 +11,13 @@ import triton.language as tl from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._utils import get_sm_version + +# Approved exception (as for `qsa_paged_sparse_gqa`): QSA's sparse path has no torch.ops +# surface; its split-K buckets and row-count PDL policy are reached as Python. +from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import ( + _qsa_pdl_enabled, + _qsa_splitk_launch_config, +) #: Base-2 online softmax: scores are scaled by log2(e) once. _LOG2_E = tl.constexpr(1.4426950408889634) @@ -570,30 +575,6 @@ def _nvfp4_merge_splitk_kernel( tl.store(output + row * output_stride_row + head * output_stride_head + dim_offsets, merged) -def _splitk_launch_config( - rows: int, num_kv_heads: int, selected_width: int -) -> tuple[int, int, int]: - """(BLOCK_N, splits, warps): the QSA split-K buckets by programs = rows x Hkv.""" - base_programs = rows * num_kv_heads - if base_programs <= 4: - block_n, target_splits, num_warps = 16, 64, 4 - elif base_programs < 32: - block_n, target_splits, num_warps = 16, 32, 4 - elif base_programs <= 256: - block_n, target_splits, num_warps = 64, 8, 2 - elif base_programs <= 512: - block_n, target_splits, num_warps = 64, 4, 2 - else: - block_n, target_splits, num_warps = 64, 1, 2 - num_tiles = triton.cdiv(selected_width, block_n) - max_useful_splits = 1 << (num_tiles.bit_length() - 1) - return block_n, min(max_useful_splits, target_splits), num_warps - - -def _pdl_enabled(rows: int) -> bool: - return rows <= 4 and os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" and get_sm_version() >= 90 - - def _sparse_gqa_nvfp4( *, q: torch.Tensor, @@ -749,11 +730,11 @@ def _sparse_gqa_nvfp4( ) return output - block_n, num_splits, num_warps = _splitk_launch_config( + block_n, num_splits, num_warps = _qsa_splitk_launch_config( rows, num_kv_heads, selected_tokens.shape[1] ) num_tiles = triton.cdiv(selected_tokens.shape[1], block_n) - use_pdl = num_splits > 1 and _pdl_enabled(rows) + use_pdl = num_splits > 1 and _qsa_pdl_enabled(rows) if num_splits == 1: partial_output = partial_lse = output else: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_rowtile_attention_cute.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_rowtile_attention_cute.py index 3b6b861108d9..03e2c925518e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_rowtile_attention_cute.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_rowtile_attention_cute.py @@ -2,10 +2,10 @@ # SPDX-License-Identifier: Apache-2.0 """The row-tile pass of `qsa_sparse_gqa_rowtile` on SM103 tensor cores (CuTe DSL, tcgen05). -Same inputs and output as that op's Triton `_rowtile_attention_kernel`: per group of -`rows_per_tile` rows of one request, the plan's packed slot list (`slots[group]`, `listed` -entries), each row's bitmap re-indexed onto it (`packed_bits`, as int32 halves) and -`group_plan[group] = (tiles, listed)`; groups with `tiles < 0` belong to the lone-row pass. +Inputs and output: per group of `rows_per_tile` rows of one request, the plan's packed +slot list (`slots[group]`, `listed` entries), each row's bitmap re-indexed onto it +(`packed_bits`, as int32 halves) and `group_plan[group] = (tiles, listed)`; groups with +`tiles < 0` belong to the lone-row pass. One CTA per (group, KV head), 128 MMA rows = `rows_per_tile` query rows x `group_size` heads, 9 warps: @@ -23,8 +23,7 @@ Tensor memory: O in columns 0-255, S/P buffers at 256 and 320, Q at 384 (512 columns). The tile's row max covers every listed entry, an upper bound of the row's own: a valid -softmax offset. P is rounded to bf16 for the PV MMA and summed in fp32 unrounded, as the -Triton kernel does. +softmax offset. P is rounded to bf16 for the PV MMA and summed in fp32 unrounded. """ from __future__ import annotations @@ -911,12 +910,10 @@ def _aligned(t: torch.Tensor, strides: int) -> bool: return t.data_ptr() % 16 == 0 and all(t.stride(i) % 8 == 0 for i in range(strides)) -def eligible( - q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, block_m: int, block_n: int -) -> bool: - """Whether a launch fits this pass: SM103, bf16 q and K/V, D = 256, the 128-row tile of - 64-token tiles, unit-stride head dimension and 16-byte-aligned rows.""" - if (block_m, block_n) != (BLOCK_M, BLOCK_N) or q.shape[2] != HEAD_DIM: +def eligible(q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor) -> bool: + """Whether a launch fits this pass: SM103, bf16 q and K/V, D = 256, unit-stride head + dimension and 16-byte-aligned rows.""" + if q.shape[2] != HEAD_DIM: return False if not (q.dtype == k_cache.dtype == v_cache.dtype == torch.bfloat16): return False @@ -933,7 +930,7 @@ def eligible( class _Launcher: - """`rowtile_attention_cute[grid](...)`: grid-indexed like the Triton kernel it stands in for; + """`rowtile_attention_cute[grid](...)`: grid-indexed like the entry's Triton launches; the grid is `(groups, kv_heads)`.""" def __getitem__(self, grid): diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py index b26d04ebb77f..7251913ffade 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py @@ -16,7 +16,6 @@ from __future__ import annotations import math -import os import cutlass import cutlass.cute as cute @@ -29,12 +28,13 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper from tensorrt_llm._torch.cute_dsl_kernels.blackwell.utils import make_ptr +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl from .qsa_paged_sparse_gqa import qsa_paged_sparse_gqa as _per_row_entry +from .qsa_rowtile_attention_cute import BLOCK_M, BLOCK_N from .qsa_rowtile_attention_cute import eligible as _cute_eligible from .qsa_rowtile_attention_cute import rowtile_attention_cute as _rowtile_attention_cute -_LOG2_E = 1.4426950408889634 _INTEGER_DTYPES = (torch.int32, torch.int64) # The per-row kernel's scan tile: the prefill bound is rounded up to it, so the # bitmap reads exactly the columns that kernel reads. @@ -43,75 +43,25 @@ _LONE_PROGRAMS = 256 # capacity: fixed: a launch width; the lone list may be any length # Bitmap words the plan pass unions and packs per loop step. _PLAN_WORDS = 16 # capacity: fixed: a loop step; the plan loops over every word - - -def _pdl_enabled() -> bool: - if os.environ.get("TRTLLM_ENABLE_PDL", "1") != "1": - return False - return torch.cuda.get_device_capability()[0] >= 9 - - -@triton.jit -def _rowtile_bits_kernel( - selected_tokens, - query_positions, - bits, - row_info, - selected_stride_row, - selected_stride_token, - bits_stride_row, - TOPK: tl.constexpr, - BLOCK_TOPK: tl.constexpr, - COMPRESS_RATIO: tl.constexpr, - TOKENS_PER_BLOCK: tl.constexpr, - PAGE_TABLE_WIDTH: tl.constexpr, - BOUND_TILE: tl.constexpr, - BLOCK_C: tl.constexpr, -): - """Row `r`: OR a bit per readable selected token into `bits[r]`; record the - highest such token and the number of columns read in `row_info[r]`.""" - row = tl.program_id(0).to(tl.int64) - visible_tokens = tl.load(query_positions + row) + 1 - selected_blocks = tl.minimum(visible_tokens // COMPRESS_RATIO, BLOCK_TOPK) - valid_tokens = selected_blocks * COMPRESS_RATIO + visible_tokens % COMPRESS_RATIO - columns_read = tl.minimum( - tl.cdiv(tl.minimum(valid_tokens, TOPK), BOUND_TILE) * BOUND_TILE, TOPK - ) - columns_read = tl.maximum(columns_read, 0).to(tl.int32) - offsets = tl.arange(0, BLOCK_C) - highest = tl.full([BLOCK_C], -1, tl.int32) - one = tl.full([BLOCK_C], 1, tl.int64) - for start in tl.range(0, columns_read, BLOCK_C): - columns = start + offsets - in_bound = columns < columns_read - tokens = tl.load( - selected_tokens + row * selected_stride_row + columns * selected_stride_token, - mask=in_bound, - other=-1, - ).to(tl.int32) - readable = in_bound & (tokens >= 0) & (tokens // TOKENS_PER_BLOCK < PAGE_TABLE_WIDTH) - safe = tl.maximum(tokens, 0) - tl.atomic_or( - bits + row * bits_stride_row + (safe >> 6), - one << (safe & 63).to(tl.int64), - mask=readable, - sem="relaxed", - ) - highest = tl.maximum(highest, tl.where(readable, tokens, -1)) - tl.store(row_info + row * 2, tl.max(highest, axis=0)) - tl.store(row_info + row * 2 + 1, columns_read) +# A one-request group takes the row-tile pass when its list's 64-slot tiles, each costing +# this many quarters of a per-row tile (2.5x), cost no more than its rows' own columns. +_SHARED_TILE_COST = 10 +_PLAN_WARPS = 4 +_LONE_WARPS = 2 +_LONE_STAGES = 2 # -------------------------------------------------------------------------------------- -# The bitmap pass in CuTe DSL (Tier 6). The Triton pass above ORs every readable selected -# token's bit into a zeroed global bitmap with a global `atomic_or` (~2051 a row). Here a -# CTA owns one row: the row's selection is staged in shared memory with coalesced loads, -# each thread walks `_BITS_RUN` consecutive columns, ORs the bits of one 64-token word in a -# register and flushes the word into a shared-memory bitmap with one `red.shared.or.b64` -# when the run leaves it; then the CTA writes every word of the row and `row_info`. The -# result is the Triton pass's, word for word, in any column order (OR is idempotent and -# commutative); the order only decides how few flushes a run makes, and the selection -# arrives in token order. No memset of the bitmap is needed. +# The bitmap pass (CuTe DSL, Tier 6). Row `r`: one bit per readable selected token (a +# column below the per-row kernel's prefill bound naming a token >= 0 whose page is inside +# the table) set in `bits[r]`, and the highest such token and the number of columns read +# in `row_info[r]`. A CTA owns one row: the row's selection is staged in shared memory with +# coalesced loads, each thread walks `_BITS_RUN` consecutive columns, ORs the bits of one +# 64-token word in a register and flushes the word into a shared-memory bitmap with one +# `red.shared.or.b64` when the run leaves it; then the CTA writes every word of the row and +# `row_info`. Any column order gives the same words (OR is idempotent and commutative); the +# order only decides how few flushes a run makes, and the selection arrives in token order. +# No memset of the bitmap is needed. # -------------------------------------------------------------------------------------- _BITS_THREADS = 128 @@ -135,7 +85,7 @@ def _red_or_shared(addr, val, *, loc=None, ip=None): class _RowBitsKernel: - """`_rowtile_bits_kernel`'s `bits` and `row_info`, one CTA a row.""" + """The bitmap pass's `bits` and `row_info`, one CTA a row.""" def __init__(self, topk, block_topk, compress_ratio, tokens_per_block, table_width, num_words): self.topk, self.block_topk, self.cr = topk, block_topk, compress_ratio @@ -245,23 +195,13 @@ def kernel(self, mSel: cute.Tensor, mQp: cute.Tensor, mBits: cute.Tensor, mInfo: def _row_bits( selected_tokens, query_positions, compress_ratio, tokens_per_block, table_width, num_words ): - """`(bits [rows, num_words] int64, row_info [rows, 2] int32)` of the bitmap pass: the - CuTe kernel when the selection has unit column stride, the Triton pass otherwise.""" + """`(bits [rows, num_words] int64, row_info [rows, 2] int32)` of the bitmap pass over a + selection with unit column stride.""" rows, final_topk = selected_tokens.shape block_topk = (final_topk - compress_ratio + 1) // compress_ratio device = selected_tokens.device row_info = torch.empty((rows, 2), dtype=torch.int32, device=device) query_positions = query_positions.contiguous() - if selected_tokens.stride(1) != 1: - bits = torch.zeros((rows, num_words), dtype=torch.int64, device=device) - _rowtile_bits_kernel[(rows,)]( - selected_tokens, query_positions, bits, row_info, - selected_tokens.stride(0), selected_tokens.stride(1), bits.stride(0), - TOPK=final_topk, BLOCK_TOPK=block_topk, COMPRESS_RATIO=compress_ratio, - TOKENS_PER_BLOCK=tokens_per_block, PAGE_TABLE_WIDTH=table_width, - BOUND_TILE=_BOUND_TILE, BLOCK_C=1024, num_warps=4, - ) # fmt: skip - return bits, row_info bits = torch.empty((rows, num_words), dtype=torch.int64, device=device) key = ( final_topk, compress_ratio, tokens_per_block, table_width, @@ -402,123 +342,6 @@ def _rowtile_plan_kernel( tl.store(lone_rows + base + offsets, members.to(tl.int32), mask=live) -@triton.jit -def _rowtile_attention_kernel( - q, k_cache, v_cache, slots, packed_bits, group_plan, output, - softmax_scale, rows, - q_stride_row: tl.constexpr, q_stride_head: tl.constexpr, q_stride_dim: tl.constexpr, - k_stride_page: tl.constexpr, k_stride_head: tl.constexpr, - k_stride_token: tl.constexpr, k_stride_dim: tl.constexpr, - v_stride_page: tl.constexpr, v_stride_head: tl.constexpr, - v_stride_token: tl.constexpr, v_stride_dim: tl.constexpr, - slots_stride_group: tl.constexpr, bits_stride_row: tl.constexpr, - output_stride_row: tl.constexpr, output_stride_head: tl.constexpr, - output_stride_dim: tl.constexpr, - GROUP_SIZE: tl.constexpr, TOKENS_PER_BLOCK: tl.constexpr, ROWS_PER_TILE: tl.constexpr, - BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, - HEAD_DIM: tl.constexpr, RESCALE_THRESHOLD: tl.constexpr, NUM_STAGES: tl.constexpr, - WARP_SPECIALIZE: tl.constexpr, USE_PDL: tl.constexpr, -): # fmt: skip - """A group of rows of one request: walk the group's packed slot list in - 64-slot tiles, each tile's K/V loaded once for all rows and heads, each - row masked to its own packed bitmap word. Groups the plan did not take - exit here and are served row by row by `_rowtile_lone_kernel`.""" - group = tl.program_id(0).to(tl.int64) - kv_head = tl.program_id(1) - tiles = tl.load(group_plan + group * 2) - if tiles < 0: - if USE_PDL: - # The lone pass may still be running: exit only once it has, so a - # later launch on the stream sees every row written. - tl.extra.cuda.gdc_wait() - return - listed = tl.load(group_plan + group * 2 + 1) - slot_base = slots + group * slots_stride_group - m = tl.arange(0, BLOCK_M) - local_row = m // GROUP_SIZE - head = m % GROUP_SIZE - row = group * ROWS_PER_TILE + local_row - live = (local_row < ROWS_PER_TILE) & (row < rows) - safe_row = tl.where(live, row, 0) - - dim_offsets = tl.arange(0, HEAD_DIM) - query_heads = kv_head * GROUP_SIZE + head - q_tile = tl.load( - q + safe_row[:, None] * q_stride_row + query_heads[:, None] * q_stride_head - + dim_offsets[None, :] * q_stride_dim, - mask=live[:, None], - other=0.0, - ) # fmt: skip - accumulator = tl.zeros([BLOCK_M, HEAD_DIM], tl.float32) - # Finite, so a row with nothing yet attended rescales by exactly zero. - running_max = tl.full([BLOCK_M], -1.0e30, tl.float32) - running_sum = tl.zeros([BLOCK_M], tl.float32) - token_offsets = tl.arange(0, BLOCK_N) - # The packed bitmap as int32 halves: word 2t holds list entries [64t, 64t+32), 2t+1 the rest. - word_base = packed_bits + safe_row * bits_stride_row - in_high_half = token_offsets >= 32 - # Moves bit (n % 32) of a word to the sign bit, so a select is one compare. - to_sign = 31 - token_offsets % 32 - for tile in tl.range(0, tiles, num_stages=NUM_STAGES, warp_specialize=WARP_SPECIALIZE): - low = tl.load(word_base + 2 * tile, mask=live, other=0) - high = tl.load(word_base + 2 * tile + 1, mask=live, other=0) - entries = tile * BLOCK_N + token_offsets - usable = entries < listed - cache_slots = tl.load(slot_base + entries, mask=usable, other=0) - safe_pages = (cache_slots // TOKENS_PER_BLOCK).to(tl.int64) - token_in_page = cache_slots % TOKENS_PER_BLOCK - # K and V tiles share one [slots, dim] layout, so the slot list is read - # in one layout per tile; the QK product takes K transposed. - keys = tl.trans(tl.load( - k_cache + safe_pages[:, None] * k_stride_page + kv_head * k_stride_head - + token_in_page[:, None] * k_stride_token + dim_offsets[None, :] * k_stride_dim, - mask=usable[:, None], - other=0.0, - ).to(q_tile.dtype)) # fmt: skip - values = tl.load( - v_cache + safe_pages[:, None] * v_stride_page + kv_head * v_stride_head - + token_in_page[:, None] * v_stride_token + dim_offsets[None, :] * v_stride_dim, - mask=usable[:, None], - other=0.0, - ).to(q_tile.dtype) # fmt: skip - # A row's bits are set only on listed entries, so they already imply residency. - word = tl.where(in_high_half[None, :], high[:, None], low[:, None]) - attend = (word << to_sign[None, :]) < 0 - scores = tl.dot(q_tile, keys) * (softmax_scale * 1.4426950408889634) - scores = tl.where(attend, scores, -float("inf")) - next_max = tl.maximum(running_max, tl.max(scores, axis=1)) - if RESCALE_THRESHOLD > 0: - # Rescale only when some row's maximum grew past the threshold; - # otherwise exponentiate against the stale maximum (probabilities - # stay below 2**RESCALE_THRESHOLD; the final division is unchanged). - if tl.max(next_max - running_max, axis=0) > RESCALE_THRESHOLD: - correction = tl.math.exp2(running_max - next_max) - accumulator = accumulator * correction[:, None] - running_sum = running_sum * correction - running_max = next_max - else: - correction = tl.math.exp2(running_max - next_max) - accumulator = accumulator * correction[:, None] - running_sum = running_sum * correction - running_max = next_max - # Unattended scores are -inf, so their probabilities are exactly zero. - probabilities = tl.math.exp2(scores - running_max[:, None]) - accumulator = tl.dot(probabilities.to(q_tile.dtype), values, accumulator) - running_sum = running_sum + tl.sum(probabilities, axis=1) - - normalized = tl.where( - running_sum[:, None] > 0, accumulator / tl.maximum(running_sum[:, None], 1.0e-20), 0.0 - ) - tl.store( - output + safe_row[:, None] * output_stride_row + query_heads[:, None] * output_stride_head - + dim_offsets[None, :] * output_stride_dim, - normalized.to(output.dtype.element_ty), - mask=live[:, None], - ) # fmt: skip - if USE_PDL: - tl.extra.cuda.gdc_wait() - - @triton.jit def _rowtile_lone_kernel( q, k_cache, v_cache, block_table, selected_tokens, request_indices, row_info, @@ -663,12 +486,12 @@ def _check_inputs(q, k_cache, v_cache, block_table, selected_tokens, request_ind raise ValueError("QSA cache token dimension must match positive tokens_per_block") if k_cache.shape[3] != head_dim or k_cache.dtype != v_cache.dtype: raise ValueError("QSA query and K/V cache head geometry is incompatible") - if q.dtype not in (torch.bfloat16, torch.float16): - raise ValueError("QSA queries must use BF16 or FP16") - if k_cache.dtype not in (torch.bfloat16, torch.float16, torch.float8_e4m3fn): - raise ValueError("QSA K/V cache must use BF16, FP16, or FP8 E4M3") - if head_dim < 16 or head_dim & (head_dim - 1): - raise ValueError("QSA sparse GQA requires a power-of-two head dimension") + if q.dtype != torch.bfloat16: + raise ValueError("QSA queries must use BF16") + if k_cache.dtype != torch.bfloat16: + raise ValueError("QSA K/V cache must use BF16") + if selected_tokens.stride(1) != 1: + raise ValueError("QSA selected tokens must have unit column stride") num_kv_heads = k_cache.shape[1] if num_kv_heads <= 0 or num_q_heads <= 0 or num_q_heads % num_kv_heads: raise ValueError("QSA query heads must be divisible by non-zero KV heads") @@ -691,6 +514,13 @@ def _check_inputs(q, k_cache, v_cache, block_table, selected_tokens, request_ind request_indices, query_positions)} # fmt: skip if len(devices) != 1 or next(iter(devices)).type != "cuda": raise ValueError("QSA row-tile sparse GQA needs every tensor on one CUDA device") + if num_q_heads // num_kv_heads > BLOCK_M: + raise ValueError("a row's query heads per KV head must fit the 128-row tile") + if not _cute_eligible(q, k_cache, v_cache): + raise ValueError( + "the row-tile pass needs SM103, head_dim 256, unit-stride heads (q's heads packed) " + "and 16-byte-aligned rows" + ) def qsa_rowtile_sparse_gqa( @@ -705,40 +535,23 @@ def qsa_rowtile_sparse_gqa( softmax_scale: float, query_positions: torch.Tensor, compress_ratio: int, - block_m: int = 128, - block_n: int = 64, - num_warps: int = 8, - num_stages: int = 3, - rescale_threshold: float = 0.0, - warp_specialize: bool = False, - shared_tile_cost: int = 10, - lone_warps: int = 2, - lone_stages: int = 2, - plan_warps: int = 4, - plan_words: int = _PLAN_WORDS, ) -> torch.Tensor: - """Bitmap pass, plan pass, lone-row pass, row-tile pass; returns `[rows, Hq, D]` contiguous. - The row-tile pass runs `qsa_rowtile_attention_cute` where it is eligible (SM103, bf16 q and - K/V, D = 256, the default tile) and the Triton kernel otherwise; the Triton launch knobs - (`num_warps`, `num_stages`, `rescale_threshold`, `warp_specialize`) apply to the latter.""" + """Bitmap pass, plan pass, lone-row pass, row-tile pass (`qsa_rowtile_attention_cute`); + returns `[rows, Hq, D]` contiguous.""" _check_inputs(q, k_cache, v_cache, block_table, selected_tokens, request_indices, tokens_per_block, softmax_scale, query_positions, compress_ratio) # fmt: skip rows, num_q_heads, head_dim = q.shape num_kv_heads = k_cache.shape[1] group_size = num_q_heads // num_kv_heads - if block_m < group_size or block_m & (block_m - 1): - raise ValueError("block_m must be a power of two holding at least one row's heads") output = torch.empty((rows, num_q_heads, head_dim), dtype=q.dtype, device=q.device) if rows == 0: return output table_width = block_table.shape[1] - if block_n != 64: - raise ValueError("block_n must be the 64-token bitmap word") num_words = triton.cdiv(table_width * tokens_per_block, 64) bits, row_info = _row_bits( selected_tokens, query_positions, compress_ratio, tokens_per_block, table_width, num_words ) - rows_per_tile = block_m // group_size + rows_per_tile = BLOCK_M // group_size rows_pow2 = triton.next_power_of_2(rows_per_tile) kv_strides = ( k_cache.stride(0), k_cache.stride(1), k_cache.stride(2), k_cache.stride(3), @@ -764,35 +577,23 @@ def qsa_rowtile_sparse_gqa( lone_rows, lone_count, rows, bits.stride(0), block_table.stride(0), block_table.stride(1), slots.stride(0), ROWS_PER_TILE=rows_per_tile, NUM_REQUESTS=block_table.shape[0], ROWS_POW2=rows_pow2, - BLOCK_N=block_n, SHARED_TILE_COST=shared_tile_cost, NUM_CACHE_PAGES=k_cache.shape[0], + BLOCK_N=BLOCK_N, SHARED_TILE_COST=_SHARED_TILE_COST, NUM_CACHE_PAGES=k_cache.shape[0], PAGE_TABLE_WIDTH=table_width, TOKENS_PER_BLOCK=tokens_per_block, - PLAN_WORDS=plan_words, num_warps=plan_warps, + PLAN_WORDS=_PLAN_WORDS, num_warps=_PLAN_WARPS, ) # fmt: skip - use_pdl = _pdl_enabled() + use_pdl = get_env_enable_pdl() _rowtile_lone_kernel[(min(rows, _LONE_PROGRAMS), num_kv_heads)]( q, k_cache, v_cache, block_table, selected_tokens, request_indices, row_info, lone_rows, lone_count, output, softmax_scale, rows, *common, selected_tokens.stride(0), selected_tokens.stride(1), *out_strides, BLOCK_H=max(16, triton.next_power_of_2(group_size)), BLOCK_N=64, - LONE_PROGRAMS=min(rows, _LONE_PROGRAMS), USE_PDL=use_pdl, LONE_STAGES=lone_stages, - num_warps=lone_warps, **geometry, + LONE_PROGRAMS=min(rows, _LONE_PROGRAMS), USE_PDL=use_pdl, LONE_STAGES=_LONE_STAGES, + num_warps=_LONE_WARPS, **geometry, ) # fmt: skip - if _cute_eligible(q, k_cache, v_cache, block_m, block_n): - _rowtile_attention_cute[(num_groups, num_kv_heads)]( - q, k_cache, v_cache, slots, packed_bits, group_plan, output, softmax_scale, rows, - group_size=group_size, rows_per_tile=rows_per_tile, tokens_per_block=tokens_per_block, - use_pdl=use_pdl, - ) # fmt: skip - return output - _rowtile_attention_kernel[(num_groups, num_kv_heads)]( - q, k_cache, v_cache, slots, packed_bits.view(torch.int32), group_plan, output, - softmax_scale, rows, *q_strides, *kv_strides, slots.stride(0), 2 * packed_bits.stride(0), - *out_strides, - GROUP_SIZE=group_size, TOKENS_PER_BLOCK=tokens_per_block, ROWS_PER_TILE=rows_per_tile, - BLOCK_M=block_m, BLOCK_N=block_n, HEAD_DIM=head_dim, - RESCALE_THRESHOLD=rescale_threshold, NUM_STAGES=num_stages, - WARP_SPECIALIZE=warp_specialize, USE_PDL=use_pdl, - num_warps=num_warps, launch_pdl=use_pdl, + _rowtile_attention_cute[(num_groups, num_kv_heads)]( + q, k_cache, v_cache, slots, packed_bits, group_plan, output, softmax_scale, rows, + group_size=group_size, rows_per_tile=rows_per_tile, tokens_per_block=tokens_per_block, + use_pdl=use_pdl, ) # fmt: skip return output @@ -858,9 +659,9 @@ class QsaSparseGqaRowtile(OpWrapper): (group, KV head), loads each listed tile's K/V once for the 120 (row, head) pairs and masks each row to its own re-indexed bits; it is launched programmatically dependent on the lone pass, whose rows are disjoint, so - the two overlap. On SM103 with bf16 q and K/V and D = 256 that pass is - `qsa_rowtile_attention_cute` (tcgen05 MMAs, S / P / O in tensor memory, - a cp.async gather of the K/V rows), elsewhere a Triton kernel. + the two overlap. That pass is `qsa_rowtile_attention_cute` (tcgen05 MMAs, + S / P / O in tensor memory, a cp.async gather of the K/V rows), so the entry + serves SM103, bf16 q and K/V and D = 256. Fusion boundary: as `qsa_paged_sparse_gqa` -- the block-table gather, scale, online softmax and P.V are inside; deciding what to attend to, @@ -913,14 +714,6 @@ class QsaSparseGqaRowtile(OpWrapper): + _decodes(23) ), ), # fmt: skip - Cell( - why="fp16 pool and fp16 q", - spec=_spec([(1900, 300)], cache_dtype=torch.float16, q_dtype=torch.float16), - ), - Cell( - why="fp8-e4m3 pool under a bf16 q, read unscaled", - spec=_spec([(1900, 300)], cache_dtype=torch.float8_e4m3fn), - ), Cell( why="int64 request indices and positions", spec=_spec([(1000, 300)], request_dtype=torch.int64), @@ -966,10 +759,9 @@ class QsaSparseGqaRowtile(OpWrapper): leave almost no gap in the union, and the call costs what the full walk does (16292-row mixed step: 3859 vs 3836 us). Below a few hundred rows it is slower than the per-row kernel (generation-only launches 1.4-2x); - the campaign's context launches are thousands of rows. The tensor-core + the target's served context launches are thousands of rows. The tensor-core row-tile pass alone, on two dumped 16384-row context launches of the - target (GB300, same plan): 631 and 887 us against 1597 and 2313 us for - the Triton pass. + target (GB300): 631 and 887 us. Scratch per call: the bitmap and the re-indexed bitmap, each `rows x ceil(table_width * tpb / 64)` int64 words (21 MB at 16384 rows and 320 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_expand_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_expand_decode.py index 42be668f474f..f28662cc23fc 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_expand_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_expand_decode.py @@ -20,7 +20,6 @@ from __future__ import annotations import functools -import os import cuda.bindings.driver as cuda_drv import cutlass @@ -31,11 +30,14 @@ from cutlass.cute.runtime import from_dlpack from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl from .qsa_expand_block_indices import qsa_expand_block_indices as _expand #: Consecutive scores a thread holds (one 16-byte load). KPT = 4 +#: Score columns a row may have: `KPT` a thread, at most 1024 threads a CTA. +MAX_WIDTH = KPT * 1024 #: CTAs an SM the register allocation is held to. At the 2592-wide score table a row is 672 #: threads, which the per-subpartition register file holds once an SM at the compiler's own #: 48 registers a thread; at three (32 a thread) rows overlap each other's barriers and @@ -53,13 +55,6 @@ _SCAN_THREADS = _HIST // 4 -@functools.lru_cache(maxsize=None) -def pdl_enabled(device_index: int) -> bool: - """Programmatic dependent launch: sm_90+ unless TRTLLM_ENABLE_PDL=0.""" - props = torch.cuda.get_device_properties(device_index) - return props.major >= 9 and os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" - - def _threads(width: int) -> int: return max(_SCAN_THREADS, -(-width // (KPT * 32)) * 32) @@ -415,7 +410,7 @@ def launch( sel = torch.empty((rows, token_topk + compress_ratio - 1), dtype=torch.int32, device=dev) if rows == 0: return sel - pdl = pdl and pdl_enabled(dev.index if dev.index is not None else torch.cuda.current_device()) + pdl = pdl and get_env_enable_pdl() fn = _compiled( rows, width, @@ -440,6 +435,11 @@ def launch( return sel +def supports(scores) -> bool: + """Whether a score table this wide is one this entry serves (see `is_valid`).""" + return scores.shape[1] % KPT == 0 and scores.shape[1] <= MAX_WIDTH + + class QsaTopkExpandDecode(OpWrapper): """`qsa_topk_select` (zero row starts, `row_ends` the visible block counts), `qsa_canonical_selection` and `qsa_expand_block_indices`, in one launch: one CTA a row, for @@ -569,7 +569,7 @@ def is_valid( assert scores.dtype == torch.float32 and scores.stride(1) == 1, ( "fp32 scores, unit column stride" ) - assert scores.shape[1] % KPT == 0 and scores.shape[1] <= KPT * 1024, "W % 4 == 0, W <= 4096" + assert scores.shape[1] % KPT == 0 and scores.shape[1] <= MAX_WIDTH, "W % 4 == 0, W <= 4096" assert scores.stride(0) % KPT == 0 and scores.data_ptr() % 16 == 0, "16-byte aligned rows" assert topk_indices.dtype == torch.int32 and topk_indices.is_contiguous(), ( "int32 topk_indices" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_attn_decode_prologue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_attn_decode_prologue.py index 97f59f172d53..8338738fc361 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_attn_decode_prologue.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_attn_decode_prologue.py @@ -21,7 +21,6 @@ from __future__ import annotations -import os from typing import Optional, Tuple import torch @@ -29,6 +28,7 @@ import triton.language as tl from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl from .qsa_paged_index_store import qsa_paged_index_store as _index_store from .qsa_paged_kv_store import qsa_paged_kv_store as _kv_store @@ -336,10 +336,6 @@ def _attn_decode_prologue_kernel( tl.store(index_cache + page * s_ic_page + slot * s_ic_slot + dims, stored_k) -def _pdl() -> bool: - return os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" - - class Qwen4ExpAttnDecodePrologue(OpWrapper): """`qwen4_exp_qk_norm_rope`, `qsa_paged_kv_store`, `qsa_paged_index_store` and `qsa_prefill_compress` of one sparse decode attention layer, in one launch. @@ -422,7 +418,7 @@ def __call__( hpb, warps = _qk_launch(num_tokens) main_programs = triton.cdiv(num_heads_q + num_heads_k, hpb) section = compress_mrope_section or (0, 0, 0) - pdl = _pdl() + pdl = get_env_enable_pdl() _attn_decode_prologue_kernel[(num_tokens, main_programs + 2)]( proj, proj.stride(0), diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_qk_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_qk_norm_rope.py index cb1a60d7c181..7e1d19411593 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_qk_norm_rope.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_qk_norm_rope.py @@ -146,8 +146,8 @@ def _qwen4_exp_qk_norm_rope_kernel( def _launch(num_tokens: int) -> Tuple[int, int]: """Main-plane heads per program and warps per program, by token count. - Measured on GB300 (microbench `qk-norm-rope-decode`, CUDA-graph replay, 12 - rotating layers, us per call; HPB/warps): + Measured on GB300 (CUDA-graph replay, 12 rotating layers, us per call; + HPB/warps): tokens 1/2 2/2 2/4 4/4 8/8 32/8 1 1.76 1.84 1.87 1.99 1.92 2.82 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/bf16_gemv_pdl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/bf16_gemv_pdl.py index 2e9d3e382f15..68163792965f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/bf16_gemv_pdl.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/bf16_gemv_pdl.py @@ -28,7 +28,6 @@ from __future__ import annotations import functools -import os import cuda.bindings.driver as cuda_drv import cutlass @@ -47,6 +46,7 @@ OpWrapper, true_fp32_matmul, ) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl #: Largest row count one launch serves; every row re-reads the weight from registers. MAX_ROWS = 4 @@ -67,13 +67,6 @@ MIN_BLOCKS_SMEM = 4 -@functools.lru_cache(maxsize=None) -def pdl_enabled(device_index: int) -> bool: - """Programmatic dependent launch: sm_90+ unless TRTLLM_ENABLE_PDL=0.""" - props = torch.cuda.get_device_properties(device_index) - return props.major >= 9 and os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" - - def tiling(n: int, k: int) -> tuple[int, int, int, int]: """`(threads, outs, souts, vec)` for a `[*, K] x [K, N]` product: register and shared-memory outputs a CTA. Without a shared-memory tier that divides N the outputs shrink to divide N.""" @@ -292,12 +285,12 @@ def _compiled(rows, n, k, pdl): def launch(mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool) -> torch.Tensor: - """The launch behind `Bf16GemvPdl.__call__`. `pdl` is honoured only where - `pdl_enabled` allows it.""" + """The launch behind `Bf16GemvPdl.__call__`. `pdl` is honoured unless + `TRTLLM_ENABLE_PDL=0`.""" m, k = mat_a.shape n = mat_b.shape[1] dev = mat_a.device - pdl = pdl and pdl_enabled(dev.index if dev.index is not None else torch.cuda.current_device()) + pdl = pdl and get_env_enable_pdl() out = torch.empty((m, n), dtype=mat_a.dtype, device=dev) if m > 0: fn = _compiled(m, n, k, bool(pdl)) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first.py deleted file mode 100644 index a307fc61006f..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first.py +++ /dev/null @@ -1,491 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Decode-row hyper-connection site, first half: the previous block's gated -injection and the mixer's first projection (`[down; block_inject]`) as one -Triton GEMV launch that leaves per-stream partial logits and per-stream sums -of squares for `qwen4_exp_hc_decode_mix`. - -Tier 5 (Triton written for this target). At decode row counts the first -projection is a GEMV over a 6.6 MB weight; cuBLAS reaches ~1.5 TB/s on it, and -the injection, the norm and the activations were three more launches. The RMS -is a per-(row, stream) scalar, so the projection runs on the *un-normed* -stream scaled by `(1 + norm_weight)`, one stream per program column: each -program's dot is exactly one stream's partial, and the consumer applies -`rstd` when it combines the four. No split-K workspace, atomics or counters. -""" - -from __future__ import annotations - -import functools -import os -from typing import Optional, Tuple - -import torch -import triton -import triton.language as tl - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -#: Largest row count this entry is certified for (the forward's dispatch threshold -#: must not exceed it). -MAX_ROWS = 16 - - -@functools.lru_cache(maxsize=None) -def pdl_enabled(device_index: int) -> bool: - """Programmatic dependent launch: sm_90+ unless TRTLLM_ENABLE_PDL=0.""" - props = torch.cuda.get_device_properties(device_index) - return props.major >= 9 and os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" - - -@triton.jit -def _hc_decode_first_kernel( - hc_ptr, - block_ptr, - inj_ptr, - shared_ptr, - gate_ptr, - topk_w_ptr, - topk_row_ptr, - norm_w_ptr, - w1_ptr, - hc_out_ptr, - part_ptr, - ss_ptr, - stride_hc, - stride_block, - stride_inj, - stride_shared, - stride_gate, - stride_topk_w, - stride_topk_row, - stride_hc_out, - stride_w1, - M: tl.constexpr, - HIDDEN: tl.constexpr, - N_REAL: tl.constexpr, - N_PAD: tl.constexpr, - BLOCK_N: tl.constexpr, - KC: tl.constexpr, - NC: tl.constexpr, - NC_P2: tl.constexpr, - HAS_INJECT: tl.constexpr, - HAS_SHARED: tl.constexpr, - TOPK: tl.constexpr, - LAUNCH_WITH_PDL: tl.constexpr, -): - # grid = (BLOCK_N-row tiles of the weight, streams). One program: BLOCK_N weight - # rows across one whole stream, as a single [BLOCK_N, NC_P2, KC] load issued - # before any arithmetic. Reductions run along the tile's last two axes, which a - # warp holds, so no CTA barrier orders a memory round trip behind another. - pid_n = tl.program_id(0) - s = tl.program_id(1) - rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - jc = tl.arange(0, NC_P2) - kc = tl.arange(0, KC) - koff = jc[:, None] * KC + kc[None, :] - kmask = (jc < NC)[:, None] & (kc >= 0)[None, :] - kcol = s * HIDDEN + koff - w = tl.load( - w1_ptr + rn[:, None, None] * stride_w1 + kcol[None, :, :], - mask=(rn < N_REAL)[:, None, None] & kmask[None, :, :], - other=0.0, - ) - gam = 1.0 + tl.load(norm_w_ptr + kcol, mask=kmask, other=0.0).to(tl.float32) - if LAUNCH_WITH_PDL: - # The weights are static, so they stream while the previous kernel finishes; - # every read of an activation and every write comes after the wait. - tl.extra.cuda.gdc_launch_dependents() - tl.extra.cuda.gdc_wait() - lead = (rn == pid_n * BLOCK_N) & (pid_n == 0) - for m in tl.static_range(M): - x = tl.load(hc_ptr + m * stride_hc + kcol, mask=kmask, other=0.0).to(tl.float32) - if HAS_INJECT: - if TOPK > 0: - # The routed experts' finalize: this row's TOPK permuted FC2 rows weighted - # by their routing weights, fp32 FMAs in slot order from zero (a slot mapped - # to -1 is skipped), stored once in the rows' dtype. - b = tl.zeros([NC_P2, KC], dtype=tl.float32) - for k in tl.static_range(TOPK): - row = tl.load(topk_row_ptr + m * stride_topk_row + k) - wk = tl.load(topk_w_ptr + m * stride_topk_w + k).to(tl.float32) - r = tl.load( - block_ptr + tl.maximum(row, 0).to(tl.int64) * stride_block + koff, - mask=kmask, - other=0.0, - ).to(tl.float32) - b = tl.where(row >= 0, tl.fma(r, wk, b), b) - b = b.to(block_ptr.dtype.element_ty).to(tl.float32) - else: - b = tl.load(block_ptr + m * stride_block + koff, mask=kmask, other=0.0).to( - tl.float32 - ) - if HAS_SHARED: - # The MoE block output, `routed + sigmoid(gate) * shared` in fp32, rounded - # once to the routed output's dtype -- the value that block used to store. - sg = tl.sigmoid(tl.load(gate_ptr + m * stride_gate).to(tl.float32)) - sh = tl.load(shared_ptr + m * stride_shared + koff, mask=kmask, other=0.0) - b = (b + sg * sh.to(tl.float32)).to(block_ptr.dtype.element_ty).to(tl.float32) - g = tl.load(inj_ptr + m * stride_inj + s).to(tl.float32) - # Rounded once, like `addcmul`: the norm and every later site read the rounded value. - xr = (x + b * g).to(hc_out_ptr.dtype.element_ty) - tl.store(hc_out_ptr + m * stride_hc_out + kcol, xr, mask=kmask & (pid_n == 0)) - x = xr.to(tl.float32) - p = tl.sum(tl.sum(w.to(tl.float32) * (x * gam)[None, :, :], axis=2), axis=1) - tl.store(part_ptr + (s * M + m) * N_PAD + rn, p) - sq = tl.sum( - tl.sum(tl.broadcast_to((x * x)[None, :, :], [BLOCK_N, NC_P2, KC]), axis=2), axis=1 - ) - tl.store(ss_ptr + s * M + m + 0 * rn, sq, mask=lead) - - -#: Launch configuration: (BLOCK_N, KC, num_warps). Swept on GB300 (see the note). -DEFAULT_CONFIG = (2, 512, 2) -#: With the MoE combine every program also reads the routed and shared rows, so -#: wider weight tiles (half the programs) win: 0.2-0.3 us per site on GB300. -COMBINE_CONFIG = (4, 512, 2) -#: With the routed experts' finalize folded in (`topk_rows`), every program -#: gathers the row's top-k FC2 rows as well. -FOLD_CONFIG = (4, 512, 2) - - -def launch( - hyper_input, - block_out, - inject, - norm_weight, - first, - hc_count, - n_cols, - shared=None, - gate_logits=None, - topk_weights=None, - topk_rows=None, - config=DEFAULT_CONFIG, -): - """The launch behind `Qwen4ExpHcDecodeFirst.__call__`, tile config exposed. - Returns `(hc_out, part, ss)`.""" - rows, width = hyper_input.shape - hidden = width // hc_count - block_n, kc, warps = config - while hidden % kc: - kc //= 2 - n_tiles = triton.cdiv(n_cols, block_n) - n_pad = n_tiles * block_n - dev, dt = hyper_input.device, hyper_input.dtype - has_inject = block_out is not None - has_shared = shared is not None - top_k = topk_rows.shape[1] if topk_rows is not None else 0 - hc_out = torch.empty((rows, width), dtype=dt, device=dev) if has_inject else hyper_input - part = torch.empty((hc_count, rows, n_pad), dtype=torch.float32, device=dev) - ss = torch.empty((hc_count, rows), dtype=torch.float32, device=dev) - pdl = pdl_enabled(dev.index if dev.index is not None else torch.cuda.current_device()) - if rows > 0: - _hc_decode_first_kernel[(n_tiles, hc_count)]( - hyper_input, - block_out if has_inject else hyper_input, - inject if has_inject else hyper_input, - shared if has_shared else hyper_input, - gate_logits if has_shared else hyper_input, - topk_weights if top_k else hyper_input, - topk_rows if top_k else hyper_input, - norm_weight, - first, - hc_out, - part, - ss, - hyper_input.stride(0), - block_out.stride(0) if has_inject else 0, - inject.stride(0) if has_inject else 0, - shared.stride(0) if has_shared else 0, - gate_logits.stride(0) if has_shared else 0, - topk_weights.stride(0) if top_k else 0, - topk_rows.stride(0) if top_k else 0, - hc_out.stride(0), - first.stride(1), - M=rows, - HIDDEN=hidden, - N_REAL=n_cols, - N_PAD=n_pad, - BLOCK_N=block_n, - KC=kc, - NC=hidden // kc, - NC_P2=triton.next_power_of_2(hidden // kc), - HAS_INJECT=has_inject, - HAS_SHARED=has_shared, - TOPK=top_k, - LAUNCH_WITH_PDL=pdl, - num_warps=warps, - launch_pdl=pdl, - ) - return hc_out, part, ss - - -class Qwen4ExpHcDecodeFirst(OpWrapper): - """With `C = hc_count`, `H = width // C`, `N = n_cols`, `first` the - `[width, >= N]` view of the load-time `[down; block_inject; 0]` weight:: - - hc_out[t, s*H+i] = hyper_input[t, s*H+i] + block_out[t, i] * inject[t, s] (rounded once) - part[s, t, n] = sum_i hc_out[t, s*H+i] * (1 + norm_weight[s*H+i]) * first[s*H+i, n] - ss[s, t] = sum_i hc_out[t, s*H+i]^2 - - Returns `(hc_out, part, ss)`: `hc_out` fresh in the activation dtype (it - *is* `hyper_input` when `block_out` / `inject` are None -- the norm-only - site); `part` fp32 `[C, T, N]` (a column slice of a padded buffer) and `ss` - fp32 `[C, T]`. The grouped RMS `rstd[t, s] = (ss[s, t] / H + eps)^-1/2` - scales stream `s`'s partial: `normed @ first[:, :N] = sum_s rstd[:, s] * - part[s]`. `qwen4_exp_hc_decode_mix` makes that combination. - - With `shared` and `gate_logits` (the site after a MoE block), `block_out` - is the routed experts' output and the block output injected above is - `block_out[t, i] + sigmoid(gate_logits[t]) * shared[t, i]` in fp32, rounded - once to `block_out`'s dtype -- bitwise what `fused_sigmoid_gate_mul_add` - stores. `gate_logits` is one pre-sigmoid scalar per row, `[T, 1]` or `[T]`. - - With `topk_weights` / `topk_rows` as well (the MoE runner left unfinalized), - `block_out` is its permuted FC2 output `[P, H]` and the routed row is the - runner's finalize, formed here: `routed[t, i] = sum_k topk_weights[t, k] * - block_out[topk_rows[t, k], i]` as fp32 FMAs in slot order from zero (a slot - whose row is -1 is skipped), rounded once to `block_out`'s dtype -- bitwise - what the runner's finalize stores. `topk_rows` is the runner's - expanded->permuted index map `[T, k]` int32, `topk_weights` its `[T, k]` - routing weights. The fold needs the combine (`shared`, `gate_logits`). - - Fusion boundary: the routed experts' finalize (optional), the - shared-expert combine (optional), the previous block's - injection, the `(1 + w)` scale and the first projection per stream, plus - the per-stream sum of squares. The routed and shared experts and the gate - projection, the RMS itself, the activations and everything after are not - inside. - """ - - ARCHS = frozenset({Arch.SM_103}) - - # spec: rows, inject (True), moe (False: shared + gate_logits too), gate - # "column" [T, 1] | "flat" [T], fold (0: top-k of the finalize folded in, - # block_out then the permuted rows), drop (False: some slots -1), n_cols - # (324), hc_count (4), hidden (2560), n_rows (384: the padded weight), - # dtype (bf16). - CELLS: tuple[Cell, ...] = tuple( - Cell(why=f"checkpoint 4x2560, 324 cols of 384, bf16, inject, {why}", spec=dict(rows=t)) - for t, why in ( - (1, "decode 1 row (con 1)"), - (2, "decode 2 rows"), - (4, "decode 4 rows"), - (8, "decode 8 rows"), - (16, "decode 16 rows (MAX_ROWS)"), - ) - ) + ( - Cell(why="MoE combine, [T, 1] gate, 1 row (con 1)", spec=dict(rows=1, moe=True)), - Cell(why="MoE combine, flat [T] gate, 4 rows", spec=dict(rows=4, moe=True, gate="flat")), - Cell(why="MoE combine, 16 rows (MAX_ROWS)", spec=dict(rows=16, moe=True)), - Cell( - why="MoE combine into the stack-final mixer: 320 cols, 1 row", - spec=dict(rows=1, moe=True, n_cols=320), - ), - Cell(why="finalize fold, top-10, 1 row (con 1)", spec=dict(rows=1, moe=True, fold=10)), - Cell( - why="finalize fold into the stack-final mixer: 320 cols, 1 row", - spec=dict(rows=1, moe=True, fold=10, n_cols=320), - ), - Cell( - why="finalize fold, 4 rows, slots mapped to -1, flat gate", - spec=dict(rows=4, moe=True, fold=10, drop=True, gate="flat"), - ), - Cell(why="finalize fold, 16 rows (MAX_ROWS)", spec=dict(rows=16, moe=True, fold=10)), - Cell( - why="norm-only site (stack start / after PLE), 1 row", spec=dict(rows=1, inject=False) - ), - Cell(why="norm-only site, 8 rows", spec=dict(rows=8, inject=False)), - Cell(why="stack-final mixer: 320 cols, 1 row", spec=dict(rows=1, n_cols=320)), - Cell(why="stack-final mixer: 320 cols, 16 rows", spec=dict(rows=16, n_cols=320)), - Cell(why="hyper_input a row slice buf[::2], 4 rows", spec=dict(rows=4, layout="row_slice")), - Cell(why="unpadded weight [width, 324], 3 rows", spec=dict(rows=3, n_rows=324)), - Cell(why="hc_count 2 at hidden 2560, 2 rows", spec=dict(rows=2, hc_count=2)), - Cell( - why="hidden 128, 36 cols (not the checkpoint), 5 rows", - spec=dict(rows=5, hidden=128, n_cols=36), - ), - Cell(why="float16 operands, 4 rows", spec=dict(rows=4, dtype=torch.float16)), - Cell(why="zero rows", spec=dict(rows=0)), - ) - - note = """ - Metadata consumed: none. Call sites: `DecodeTarget`'s trunk, every - hyper-connection site (97 per forward) when the decode batch has at most - `_HC_DECODE_GEMV_MAX_ROWS` rows (<= MAX_ROWS = 16), followed by - `qwen4_exp_hc_decode_mix`. Larger batches and `PrefillTarget` keep - `qwen4_exp_hc_inject_norm` + `cublas_mm` + `qwen4_exp_hc_lowrank_act`. - The 47 sites after a MoE block other than the one before the PLE add pass - `shared` and `gate_logits`, and with them the unfinalized runner's - `topk_weights` / `topk_rows` (`block_out` its permuted FC2 rows). - - Preconditions: 2-D operands with unit column stride (row strides free); - `first` `[width, >= n_cols]` with unit stride along `width` (the `.t()` - view of a row-major `[N, width]` weight); `hidden` a multiple of 16. Row - count is a compile-time constant (one Triton specialization per decode - batch size); at most MAX_ROWS rows certified. Columns past `n_cols` are - never read, so the zero padding of the load-time weight costs no traffic. - - Numerics: `hc_out` as `qwen4_exp_hc_inject_norm` (fp32 accumulate, one - rounding); the projection multiplies the fp32 `hc_out * (1 + w)` by the - 16-bit weight in fp32 (the composition it replaces rounded the normed - stream to bf16 first); deterministic (no atomics, fixed reduction order). - The optional combine is fp32 rounded once to `block_out`'s dtype, which - against the float64 reference can land one ulp of the block output away - (its rounding midpoint); `hc_out` carries that scaled by `inject`, and the - test checks `part` / `ss` against the projection of the `hc_out` the op - stored. Every program recomputes the combine for its stream's columns: the - extra `[hidden]` read per program is L2-resident. The optional finalize is - the runner's own arithmetic (fp32 FMAs in slot order from zero, one - rounding), so `hc_out` and `ss` are bitwise the finalize + combine - composition at the same tile, and so is `part` from 2 rows; at 1 row the - fold's Triton specialization reduces the dot in another fp32 order (~12% of - `part` one fp32 ulp apart, < 1e-7 relative; GB300). The - float64 reference rounds the finalize from an exact sum, which can land one - ulp of the routed row away, and the test's band carries that term too. - With the fold every program also gathers the row's top-k FC2 rows (`k x - [hidden]`, L2-resident: the runner's GEMM wrote them just before). - Graph-capturable: no host sync, outputs allocated per call. Launched with - programmatic dependent launch on sm_90+ (`TRTLLM_ENABLE_PDL=0` turns it - off): the static weight tile streams before `gdc_wait`, every activation - read and every write after it, and dependents are released at once so - `qwen4_exp_hc_decode_mix` can stream its weights meanwhile. - - Grid: (cdiv(n_cols, 2), hc_count) programs of 2 warps, each one load of 2 - weight rows x one stream (10 KB); with the combine, 4 weight rows - (`COMBINE_CONFIG`; it sums `part` in a different fp32 order than the - default tile, inside the same band). Measured on GB300 (graph replay, cold - weights, with `qwen4_exp_hc_decode_mix`): 1 row is fast; each further row - costs several us (the row loop serializes its memory round trips), so the - forward dispatches only 1-row batches here. Hard limits: int32 offsets, far below 2^31 at <= 16 rows. - """ - - def __call__( - self, - hyper_input: torch.Tensor, - block_out: Optional[torch.Tensor], - inject: Optional[torch.Tensor], - norm_weight: torch.Tensor, - first: torch.Tensor, - hc_count: int, - n_cols: int, - shared: Optional[torch.Tensor] = None, - gate_logits: Optional[torch.Tensor] = None, - topk_weights: Optional[torch.Tensor] = None, - topk_rows: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - if topk_rows is not None: - config = FOLD_CONFIG - elif shared is not None: - config = COMBINE_CONFIG - else: - config = DEFAULT_CONFIG - hc_out, part, ss = launch( - hyper_input, - block_out, - inject, - norm_weight, - first, - hc_count, - n_cols, - shared, - gate_logits, - topk_weights, - topk_rows, - config=config, - ) - return hc_out, part[:, :, :n_cols], ss - - def reference( - self, - hyper_input: torch.Tensor, - block_out: Optional[torch.Tensor], - inject: Optional[torch.Tensor], - norm_weight: torch.Tensor, - first: torch.Tensor, - hc_count: int, - n_cols: int, - shared: Optional[torch.Tensor] = None, - gate_logits: Optional[torch.Tensor] = None, - topk_weights: Optional[torch.Tensor] = None, - topk_rows: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """float64; the finalize (with `topk_rows`) rounded once to - `block_out`'s dtype, the combine (with `shared`) rounded once to `block_out`'s - dtype, `hc_out` rounded once to the activation dtype, `part` and `ss` - rounded once to fp32.""" - dtype = hyper_input.dtype - rows, width = hyper_input.shape - hidden = width // hc_count - if block_out is None: - hc_out = hyper_input - else: - if topk_rows is not None: - valid = (topk_rows >= 0).unsqueeze(-1) - rows_ = block_out.double()[topk_rows.clamp(min=0).long()] - w = topk_weights.double().unsqueeze(-1) - block_out = torch.where(valid, rows_ * w, 0.0).sum(1).to(block_out.dtype) - if shared is not None: - sg = torch.sigmoid(gate_logits.double().reshape(rows, 1)) - block_out = (block_out.double() + sg * shared.double()).to(block_out.dtype) - streams = hyper_input.double().unflatten(-1, (hc_count, hidden)) - gated = block_out.double().unsqueeze(-2) * inject.double().unsqueeze(-1) - hc_out = (streams + gated).flatten(-2).to(dtype) - x = hc_out.double().unflatten(-1, (hc_count, hidden)) - y = x * (1.0 + norm_weight.double()).unflatten(-1, (hc_count, hidden)) - w = first[:, :n_cols].double().unflatten(0, (hc_count, hidden)) - part = torch.einsum("tsi,sin->stn", y, w).float() - ss = x.pow(2).sum(dim=-1).t().float() - return hc_out, part, ss - - def is_valid( - self, - hyper_input: torch.Tensor, - block_out: Optional[torch.Tensor], - inject: Optional[torch.Tensor], - norm_weight: torch.Tensor, - first: torch.Tensor, - hc_count: int, - n_cols: int, - shared: Optional[torch.Tensor] = None, - gate_logits: Optional[torch.Tensor] = None, - topk_weights: Optional[torch.Tensor] = None, - topk_rows: Optional[torch.Tensor] = None, - ) -> None: - # Raw row strides with unit column stride everywhere; a contiguous - # [width, N] weight would be read with its rows and columns swapped. - assert hyper_input.dim() == 2 and hyper_input.stride(1) == 1, "hyper_input [T, width]" - rows, width = hyper_input.shape - assert width % hc_count == 0 and (width // hc_count) % 16 == 0, "hidden multiple of 16" - assert (block_out is None) == (inject is None), "block_out and inject come together" - assert (shared is None) == (gate_logits is None), "shared and gate_logits come together" - assert shared is None or block_out is not None, "the combine needs block_out" - assert (topk_weights is None) == (topk_rows is None), ( - "topk_weights and topk_rows come together" - ) - assert topk_rows is None or shared is not None, "the finalize fold needs the combine" - if block_out is not None: - if topk_rows is None: - assert block_out.shape == (rows, width // hc_count) and block_out.stride(1) == 1 - else: - # The permuted FC2 rows: any row count, read through topk_rows. - assert block_out.dim() == 2 and block_out.shape[1] == width // hc_count - assert block_out.stride(1) == 1, "permuted rows [P, hidden]" - assert topk_rows.dtype == torch.int32, "topk_rows int32 (the runner's index map)" - assert topk_rows.dim() == 2 and topk_rows.shape[0] == rows, "topk_rows [T, k]" - assert topk_rows.stride(1) == 1 and topk_weights.stride(1) == 1 - assert topk_weights.shape == topk_rows.shape, "topk_weights [T, k]" - assert inject.shape == (rows, hc_count) and inject.stride(1) == 1 - if shared is not None: - assert shared.shape == (rows, width // hc_count), "shared [T, hidden]" - assert shared.stride(1) == 1, "shared [T, hidden]" - assert shared.dtype == block_out.dtype, "shared in block_out's dtype" - assert gate_logits.shape in ((rows,), (rows, 1)), "gate_logits [T, 1] or [T]" - assert rows <= 1 or gate_logits.stride(0) != 0, "one gate logit per row" - assert norm_weight.shape == (width,) and norm_weight.is_contiguous(), "norm_weight [width]" - assert first.shape[0] == width and first.shape[1] >= n_cols, "first [width, >= n_cols]" - assert first.stride(0) == 1, "first must be the .t() view of a row-major [N, width] weight" - - -qwen4_exp_hc_decode_first = Qwen4ExpHcDecodeFirst() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first_cute.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first_cute.py index 8d28cf523702..6e813be14845 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first_cute.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first_cute.py @@ -4,13 +4,11 @@ block's injection (with the MoE block's combine and the routed experts' finalize folded in, when given) and the first projection per stream, one PDL launch. -Tier 5 (CuTe DSL written for this target). `qwen4_exp_hc_decode_first` computes the -same thing for up to 16 rows in Triton, whose power-of-two tiles pad a -2560-wide stream to 8 x 512: three eighths of every program's weight, activation -and scale tiles are masked lanes holding registers. Here a stream is exactly -`THREADS x VEC` elements, one 16-byte vector per thread, and the grid is sized so -that every CTA is resident at once (one wave), so all of them stream their weight -before the producer finishes. +Tier 5 (CuTe DSL written for this target). A stream is exactly `THREADS x VEC` +elements, one 16-byte vector per thread (a power-of-two tile would pad a 2560-wide +stream to 8 x 512, three eighths of it masked lanes holding registers), and the grid +is sized so that every CTA is resident at once (one wave), so all of them stream their +weight before the producer finishes. Order of work inside a CTA (`OUTS` weight rows x one stream): @@ -28,7 +26,6 @@ from __future__ import annotations import functools -import os from typing import Optional, Tuple import cuda.bindings.driver as cuda_drv @@ -40,11 +37,10 @@ from cutlass.cute.runtime import from_dlpack from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl -from .qwen4_exp_hc_decode_first import qwen4_exp_hc_decode_first as _triton_first - -#: One decode row per launch (the 1-row batch is the one the Triton entry is slowest at -#: relative to its weight stream). +#: One decode row per launch: the forward runs only 1-row decode batches through the +#: site's GEMV pair (`_HC_DECODE_GEMV_MAX_ROWS`). MAX_ROWS = 1 VEC = 8 #: Threads per CTA: one 16-byte vector of a 2560-wide stream each. @@ -54,13 +50,6 @@ _COMPILE_OPTIONS = "--ptxas-options -maxrregcount=128" -@functools.lru_cache(maxsize=None) -def pdl_enabled(device_index: int) -> bool: - """Programmatic dependent launch: sm_90+ unless TRTLLM_ENABLE_PDL=0.""" - props = torch.cuda.get_device_properties(device_index) - return props.major >= 9 and os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" - - def outs_for(n_cols: int, hc_count: int) -> int: """Weight rows per CTA: the fewest that divide `n_cols` and keep the grid in one wave.""" for outs in range(1, n_cols + 1): @@ -354,7 +343,7 @@ def launch( ss = torch.empty((hc_count, rows), dtype=torch.float32, device=dev) if rows == 0: return hc_out, part, ss - pdl = pdl_enabled(dev.index if dev.index is not None else torch.cuda.current_device()) + pdl = get_env_enable_pdl() weight = first.t() # the row-major [n_rows, width] weight gate = gate_logits.reshape(1) if has_shared else hyper_input.new_empty(1) n_perm = block_out.shape[0] if (has_inject and top_k) else 1 @@ -396,23 +385,49 @@ def launch( class Qwen4ExpHcDecodeFirstCute(OpWrapper): - """`qwen4_exp_hc_decode_first`'s computation and return values for a one-row batch - (see that entry for the formulas, the MoE combine and the finalize fold):: + """The first half of a one-row decode hyper-connection site: the previous block's + injection and the mixer's first projection (`[down; block_inject]`) per stream. With + `C = hc_count`, `H = width // C`, `N = n_cols`, `first` the `[width, >= N]` view of + the load-time `[down; block_inject; 0]` weight:: hc_out[0, s*H+i] = hyper_input[0, s*H+i] + block_out[0, i] * inject[0, s] (rounded once) part[s, 0, n] = sum_i hc_out[0, s*H+i] * (1 + norm_weight[s*H+i]) * first[s*H+i, n] ss[s, 0] = sum_i hc_out[0, s*H+i]^2 - Returns `(hc_out, part, ss)`: `hc_out` fresh (it *is* `hyper_input` at a norm-only - site), `part` fp32 `[C, 1, N]`, `ss` fp32 `[C, 1]`. The combine (`shared`, - `gate_logits`) and the fold (`topk_weights`, `topk_rows`) are the Triton entry's. - - Fusion boundary: as `qwen4_exp_hc_decode_first`. + Returns `(hc_out, part, ss)`: `hc_out` fresh in the activation dtype (it *is* + `hyper_input` when `block_out` / `inject` are None -- the norm-only site), `part` + fp32 `[C, 1, N]`, `ss` fp32 `[C, 1]`. The grouped RMS `rstd[s] = (ss[s, 0] / H + + eps)^-1/2` scales stream `s`'s partial: `normed @ first[:, :N] = sum_s rstd[s] * + part[s]`. `qwen4_exp_hc_decode_mix_cute` makes that combination. + + With `shared` and `gate_logits` (the site after a MoE block), `block_out` is the + routed experts' output and the block output injected above is `block_out[0, i] + + sigmoid(gate_logits[0]) * shared[0, i]` in fp32, rounded once to `block_out`'s + dtype -- bitwise what `fused_sigmoid_gate_mul_add` stores. `gate_logits` is one + pre-sigmoid scalar per row, `[T, 1]` or `[T]`. + + With `topk_weights` / `topk_rows` as well (the MoE runner left unfinalized), + `block_out` is its permuted FC2 output `[P, H]` and the routed row is the runner's + finalize, formed here: `routed[0, i] = sum_k topk_weights[0, k] * + block_out[topk_rows[0, k], i]` as fp32 FMAs in slot order from zero (a slot whose + row is -1 is skipped), rounded once to `block_out`'s dtype -- bitwise what the + runner's finalize stores. `topk_rows` is the runner's expanded->permuted index map + `[T, k]` int32, `topk_weights` its `[T, k]` routing weights. The fold needs the + combine (`shared`, `gate_logits`). + + Fusion boundary: the routed experts' finalize (optional), the shared-expert combine + (optional), the previous block's injection, the `(1 + w)` scale and the first + projection per stream, plus the per-stream sum of squares. The routed and shared + experts and the gate projection, the RMS itself, the activations and everything + after are not inside. """ ARCHS = frozenset({Arch.SM_103}) - # spec: the Triton entry's (rows is 1). + # spec: rows (1), inject (True), moe (False: shared + gate_logits too), gate + # "column" [T, 1] | "flat" [T], fold (0: top-k of the finalize folded in, block_out + # then the permuted rows), drop (False: some slots -1), n_cols (324), n_rows (384: + # the padded weight), layout ("row_slice": hyper_input a row slice). CELLS: tuple[Cell, ...] = ( Cell(why="checkpoint 4x2560, 324 cols of 384, bf16, inject, 1 row", spec=dict(rows=1)), Cell(why="MoE combine, [T, 1] gate, 1 row", spec=dict(rows=1, moe=True)), @@ -439,16 +454,22 @@ class Qwen4ExpHcDecodeFirstCute(OpWrapper): note = """ Metadata consumed: none. Call sites: `DecodeTarget`'s trunk, every hyper-connection site of a one-row decode step (97 per forward), followed by - `qwen4_exp_hc_decode_mix`; batches of more rows keep `qwen4_exp_hc_decode_first`. - - Preconditions: as the Triton entry, and: one row; `hidden` 2560 (320 threads x - 8); bf16 operands; `n_cols` with a divisor that keeps `n_cols / outs x C <= 152` + `qwen4_exp_hc_decode_mix_cute` (or its NVFP4 variant at an MoE site); batches of + more rows take `qwen4_exp_hc_inject_norm` and the mixer. The 47 sites after a MoE + block other than the one before the PLE add pass `shared` and `gate_logits`, and + with them the unfinalized runner's `topk_weights` / `topk_rows`. + + Preconditions: 2-D operands with unit column stride (row strides free); `first` + `[width, >= n_cols]` with unit stride along `width` (the `.t()` view of a row-major + `[N, width]` weight; columns past `n_cols` are never read); one row; `hidden` 2560 + (320 threads x 8); bf16 operands; `n_cols` with a divisor that keeps `n_cols / outs x C <= 152` CTAs (324 -> 9 rows x 36 tiles x 4 streams = 144; 320 -> 10 rows, 128); the permuted rows' count is part of the compiled layout (one compile per value). - Numerics: `hc_out` the Triton entry's arithmetic (fp32 FMAs, one rounding; the - finalize and the combine likewise), so bitwise equal to it; `part` and `ss` - reduce in another fp32 order (inside the same band of the float64 reference). + Numerics: `hc_out` as `qwen4_exp_hc_inject_norm` (fp32 FMAs, one rounding; the + finalize and the combine likewise); `part` and `ss` are fp32 reductions in a fixed + order, inside the float64 reference's band (the combine and the finalize can each + land one ulp of the block output from it, scaled into `hc_out` by `inject`). Deterministic. Graph-capturable: compiled on the first call per variant, which must precede capture (the engine's warmup forward does it). @@ -485,9 +506,47 @@ def __call__( topk_rows, ) - def reference(self, *args, **kwargs) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """The Triton entry's float64 reference.""" - return _triton_first.reference(*args, **kwargs) + def reference( + self, + hyper_input: torch.Tensor, + block_out: Optional[torch.Tensor], + inject: Optional[torch.Tensor], + norm_weight: torch.Tensor, + first: torch.Tensor, + hc_count: int, + n_cols: int, + shared: Optional[torch.Tensor] = None, + gate_logits: Optional[torch.Tensor] = None, + topk_weights: Optional[torch.Tensor] = None, + topk_rows: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """float64; the finalize (with `topk_rows`) rounded once to + `block_out`'s dtype, the combine (with `shared`) rounded once to `block_out`'s + dtype, `hc_out` rounded once to the activation dtype, `part` and `ss` + rounded once to fp32.""" + dtype = hyper_input.dtype + rows, width = hyper_input.shape + hidden = width // hc_count + if block_out is None: + hc_out = hyper_input + else: + if topk_rows is not None: + valid = (topk_rows >= 0).unsqueeze(-1) + rows_ = block_out.double()[topk_rows.clamp(min=0).long()] + w = topk_weights.double().unsqueeze(-1) + block_out = torch.where(valid, rows_ * w, 0.0).sum(1).to(block_out.dtype) + if shared is not None: + sg = torch.sigmoid(gate_logits.double().reshape(rows, 1)) + block_out = (block_out.double() + sg * shared.double()).to(block_out.dtype) + streams = hyper_input.double().unflatten(-1, (hc_count, hidden)) + gated = block_out.double().unsqueeze(-2) * inject.double().unsqueeze(-1) + hc_out = (streams + gated).flatten(-2).to(dtype) + x = hc_out.double().unflatten(-1, (hc_count, hidden)) + y = x * (1.0 + norm_weight.double()).unflatten(-1, (hc_count, hidden)) + w = first[:, :n_cols].double().unflatten(0, (hc_count, hidden)) + part = torch.einsum("tsi,sin->stn", y, w).float() + ss = x.pow(2).sum(dim=-1).t().float() + return hc_out, part, ss def is_valid( self, @@ -503,19 +562,39 @@ def is_valid( topk_weights: Optional[torch.Tensor] = None, topk_rows: Optional[torch.Tensor] = None, ) -> None: - _triton_first.is_valid( - hyper_input, - block_out, - inject, - norm_weight, - first, - hc_count, - n_cols, - shared, - gate_logits, - topk_weights, - topk_rows, + # Raw row strides with unit column stride everywhere; a contiguous + # [width, N] weight would be read with its rows and columns swapped. + assert hyper_input.dim() == 2 and hyper_input.stride(1) == 1, "hyper_input [T, width]" + rows, width = hyper_input.shape + assert width % hc_count == 0 and (width // hc_count) % 16 == 0, "hidden multiple of 16" + assert (block_out is None) == (inject is None), "block_out and inject come together" + assert (shared is None) == (gate_logits is None), "shared and gate_logits come together" + assert shared is None or block_out is not None, "the combine needs block_out" + assert (topk_weights is None) == (topk_rows is None), ( + "topk_weights and topk_rows come together" ) + assert topk_rows is None or shared is not None, "the finalize fold needs the combine" + if block_out is not None: + if topk_rows is None: + assert block_out.shape == (rows, width // hc_count) and block_out.stride(1) == 1 + else: + # The permuted FC2 rows: any row count, read through topk_rows. + assert block_out.dim() == 2 and block_out.shape[1] == width // hc_count + assert block_out.stride(1) == 1, "permuted rows [P, hidden]" + assert topk_rows.dtype == torch.int32, "topk_rows int32 (the runner's index map)" + assert topk_rows.dim() == 2 and topk_rows.shape[0] == rows, "topk_rows [T, k]" + assert topk_rows.stride(1) == 1 and topk_weights.stride(1) == 1 + assert topk_weights.shape == topk_rows.shape, "topk_weights [T, k]" + assert inject.shape == (rows, hc_count) and inject.stride(1) == 1 + if shared is not None: + assert shared.shape == (rows, width // hc_count), "shared [T, hidden]" + assert shared.stride(1) == 1, "shared [T, hidden]" + assert shared.dtype == block_out.dtype, "shared in block_out's dtype" + assert gate_logits.shape in ((rows,), (rows, 1)), "gate_logits [T, 1] or [T]" + assert rows <= 1 or gate_logits.stride(0) != 0, "one gate logit per row" + assert norm_weight.shape == (width,) and norm_weight.is_contiguous(), "norm_weight [width]" + assert first.shape[0] == width and first.shape[1] >= n_cols, "first [width, >= n_cols]" + assert first.stride(0) == 1, "first must be the .t() view of a row-major [N, width] weight" assert hyper_input.shape[0] <= MAX_ROWS, "one row" assert hyper_input.shape[1] // hc_count == THREADS * VEC, "hidden 2560" assert hyper_input.dtype == torch.bfloat16, "bf16" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix.py deleted file mode 100644 index 144ceea58438..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix.py +++ /dev/null @@ -1,331 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Decode-row hyper-connection site, second half: the grouped RMS, the SiLU -bottleneck and the injection weights from `qwen4_exp_hc_decode_first`'s -per-stream partials, then the up projection, the sigmoid gate on the normed -stream and the mean over streams -- one Triton GEMV launch. - -Tier 5 (Triton written for this target). Each program owns `BLOCK_H` hidden -columns across all streams: it requests its whole up-projection tile before -any arithmetic, rebuilds `rstd` and `low` from the four partials (a few KB of -L2 per row, recomputed by every program), and keeps the `[T, C*H]` gate -logits in registers. Program 0 also writes the block-injection weights. -""" - -from __future__ import annotations - -import functools -import os -from typing import Optional, Tuple - -import torch -import torch.nn.functional as F -import triton -import triton.language as tl - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -#: Largest row count this entry is certified for. -MAX_ROWS = 16 - - -@functools.lru_cache(maxsize=None) -def pdl_enabled(device_index: int) -> bool: - """Programmatic dependent launch: sm_90+ unless TRTLLM_ENABLE_PDL=0.""" - props = torch.cuda.get_device_properties(device_index) - return props.major >= 9 and os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" - - -@triton.jit -def _hc_decode_mix_kernel( - part_ptr, - ss_ptr, - up_ptr, - x_ptr, - norm_w_ptr, - mixed_ptr, - inj_out_ptr, - stride_ps, - stride_pm, - stride_ss, - stride_up_r, - stride_up_n, - stride_x, - stride_mixed, - eps, - M: tl.constexpr, - HIDDEN: tl.constexpr, - HC: tl.constexpr, - R1: tl.constexpr, - R2: tl.constexpr, - BLOCK_H: tl.constexpr, - WITH_INJECT_OUT: tl.constexpr, - LAUNCH_WITH_PDL: tl.constexpr, - EARLY_TRIGGER: tl.constexpr, -): - pid = tl.program_id(0) - hh = pid * BLOCK_H + tl.arange(0, BLOCK_H) - sidx = tl.arange(0, HC) - col = sidx[None, :] * HIDDEN + hh[:, None] - hmask = (hh < HIDDEN)[:, None] & (sidx >= 0)[None, :] - # The whole [BLOCK_H, HC, R] weight tile is requested before any arithmetic. - ra = tl.arange(0, R1) - w_a = tl.load( - up_ptr + col[:, :, None] * stride_up_n + ra[None, None, :] * stride_up_r, - mask=hmask[:, :, None], - other=0.0, - ) - if R2 > 0: - rb = R1 + tl.arange(0, R2) - w_b = tl.load( - up_ptr + col[:, :, None] * stride_up_n + rb[None, None, :] * stride_up_r, - mask=hmask[:, :, None], - other=0.0, - ) - gam = 1.0 + tl.load(norm_w_ptr + col, mask=hmask, other=0.0).to(tl.float32) - if LAUNCH_WITH_PDL: - if EARLY_TRIGGER: - tl.extra.cuda.gdc_launch_dependents() - # Weights above stream while the producer runs; its outputs are read after this. - tl.extra.cuda.gdc_wait() - for m in tl.static_range(M): - ssm = tl.load(ss_ptr + sidx * stride_ss + m) - pa = tl.load(part_ptr + sidx[:, None] * stride_ps + m * stride_pm + ra[None, :]) - if R2 > 0: - pb = tl.load(part_ptr + sidx[:, None] * stride_ps + m * stride_pm + rb[None, :]) - x = tl.load(x_ptr + m * stride_x + col, mask=hmask, other=0.0).to(tl.float32) - rstd = tl.rsqrt(ssm / HIDDEN + eps) - za = tl.sum(pa * rstd[:, None], axis=0) * (1.0 / HC) - acc = tl.sum(w_a.to(tl.float32) * (za * tl.sigmoid(za))[None, None, :], axis=2) - if R2 > 0: - zb = tl.sum(pb * rstd[:, None], axis=0) * (1.0 / HC) - acc += tl.sum(w_b.to(tl.float32) * (zb * tl.sigmoid(zb))[None, None, :], axis=2) - v = tl.sigmoid(acc) * (x * rstd[None, :] * gam) - tl.store( - mixed_ptr + m * stride_mixed + hh, - (tl.sum(v, axis=1) * (1.0 / HC)).to(mixed_ptr.dtype.element_ty), - mask=hh < HIDDEN, - ) - if WITH_INJECT_OUT: - ci = tl.arange(0, HC) - pi = tl.load( - part_ptr + sidx[:, None] * stride_ps + m * stride_pm + (R1 + R2) + ci[None, :], - mask=(sidx >= 0)[:, None] & (pid == 0), - other=0.0, - ) - zi = tl.sum(pi * rstd[:, None], axis=0) * (1.0 / HC) - tl.store( - inj_out_ptr + m * HC + ci, - (2.0 * tl.sigmoid(zi)).to(inj_out_ptr.dtype.element_ty), - mask=(ci >= 0) & (pid == 0), - ) - - -#: Launch configuration: (BLOCK_H, num_warps). Swept on GB300 (see the note). -DEFAULT_CONFIG = (8, 8) -#: Let the kernel after this one launch as soon as every program has requested its -#: weights: a programmatic dependent (the decode projections) then streams its own -#: weight while this kernel finishes. Its `gdc_wait` still waits for this whole grid. -TRIGGER_DEPENDENTS = True - - -def _rank_split(rank: int) -> Tuple[int, int]: - """`rank = R1 + R2` with `R1` the largest power of two not above it; `R2` must be - a power of two or zero (320 = 256 + 64).""" - r1 = 1 << (rank.bit_length() - 1) - return r1, rank - r1 - - -def launch( - part, - ss, - up, - hc_out, - norm_weight, - eps, - hc_count, - with_inject, - config=DEFAULT_CONFIG, - trigger=TRIGGER_DEPENDENTS, -): - """The launch behind `Qwen4ExpHcDecodeMix.__call__`, tile config exposed. - Returns `(mixed, inject_out)`.""" - rows, width = hc_out.shape - hidden = width // hc_count - rank = up.shape[0] - r1, r2 = _rank_split(rank) - block_h, warps = config - dev, dt = hc_out.device, hc_out.dtype - mixed = torch.empty((rows, hidden), dtype=dt, device=dev) - inject_out = torch.empty((rows, hc_count), dtype=dt, device=dev) if with_inject else None - pdl = pdl_enabled(dev.index if dev.index is not None else torch.cuda.current_device()) - if rows > 0: - _hc_decode_mix_kernel[(triton.cdiv(hidden, block_h),)]( - part, - ss, - up, - hc_out, - norm_weight, - mixed, - inject_out if with_inject else mixed, - part.stride(0), - part.stride(1), - ss.stride(0), - up.stride(0), - up.stride(1), - hc_out.stride(0), - mixed.stride(0), - eps, - M=rows, - HIDDEN=hidden, - HC=hc_count, - R1=r1, - R2=r2, - BLOCK_H=block_h, - WITH_INJECT_OUT=with_inject, - LAUNCH_WITH_PDL=pdl, - EARLY_TRIGGER=trigger, - num_warps=warps, - launch_pdl=pdl, - ) - return mixed, inject_out - - -class Qwen4ExpHcDecodeMix(OpWrapper): - """With `C = hc_count`, `H = width // C`, `R = up.shape[0]`, `part` / `ss` the - per-stream partial logits `[C, T, >= R (+ C)]` and sums of squares `[C, T]` - from `qwen4_exp_hc_decode_first`, `up` the `[R, width]` view of the - `[width, R]` up-projection weight:: - - rstd[t, s] = 1 / sqrt(ss[s, t] / H + eps) - logits[t, n] = sum_s rstd[t, s] * part[s, t, n] - low = silu(logits[:, :R] / C) - gate = sigmoid(low @ up) # [T, C*H] - normed[t, s*H+i] = hc_out[t, s*H+i] * rstd[t, s] * (1 + norm_weight[s*H+i]) - mixed[t, i] = mean over s of gate[t, s*H+i] * normed[t, s*H+i] - inject_out = 2 * sigmoid(logits[:, R:R+C] / C) # with_inject only - - fp32 inside, one rounding per output. Returns fresh `(mixed, inject_out)` - in `hc_out`'s dtype; `inject_out` is None without `with_inject`. - - Fusion boundary: the grouped RMS, the bottleneck SiLU, the injection - weights, the up projection, the sigmoid gate, the normed stream and the - stream mean. The partials are the producer's. - """ - - ARCHS = frozenset({Arch.SM_103}) - - # spec: rows, with_inject (True), hc_count (4), hidden (2560), rank (320), dtype (bf16). - CELLS: tuple[Cell, ...] = tuple( - Cell(why=f"checkpoint 4x2560 r320 bf16, inject, {why}", spec=dict(rows=t)) - for t, why in ( - (1, "decode 1 row (con 1)"), - (2, "decode 2 rows"), - (3, "decode 3 rows"), - (8, "decode 8 rows"), - (16, "decode 16 rows (MAX_ROWS)"), - ) - ) + ( - Cell(why="stack-final mixer (no inject), 1 row", spec=dict(rows=1, with_inject=False)), - Cell(why="stack-final mixer (no inject), 16 rows", spec=dict(rows=16, with_inject=False)), - Cell(why="hc_out a row slice buf[::2], 4 rows", spec=dict(rows=4, layout="row_slice")), - Cell(why="hc_count 2 at hidden 2560, 2 rows", spec=dict(rows=2, hc_count=2)), - Cell( - why="hidden 128 rank 32 (not the checkpoint), 5 rows", - spec=dict(rows=5, hidden=128, rank=32), - ), - Cell(why="float16 operands, 4 rows", spec=dict(rows=4, dtype=torch.float16)), - Cell(why="zero rows", spec=dict(rows=0)), - ) - - note = """ - Metadata consumed: none. Call sites: `DecodeTarget`'s trunk, right after - each `qwen4_exp_hc_decode_first` (97 per forward, 1-row decode batches). - - Preconditions: `part` fp32 `[C, T, >= R (+ C)]` and `ss` fp32 `[C, T]`, - both with unit last stride (other strides free); `up` `[R, width]` (any - strides; the `.t()` view of the stored weight reads rows of `R` contiguous - elements); `R` a power of two or a sum of two (320 = 256 + 64); `hc_out` - `[T, width]` unit column stride; `norm_weight` contiguous `[width]`. Row - count is a compile-time constant, at most MAX_ROWS certified. - - Numerics: fp32 throughout (the composition it replaces rounded `low`, the - gate logits and `normed` to bf16), one rounding per output. Every program - recomputes `rstd` and `low` from the partials in the same order, so all - programs agree bitwise. Graph-capturable. Launched with programmatic - dependent launch on sm_90+ (`TRTLLM_ENABLE_PDL=0` turns it off): the - weight tile streams before `gdc_wait`, the producer's outputs after it. - - Grid: cdiv(H, 8) programs of 8 warps, each one [8, C, R] weight load - (20 KB). Measured on GB300 (graph replay, cold weights): ~4 us at 1 row; - each further row costs ~3 us (the row loop serializes its L2 round trips). - """ - - def __call__( - self, - part: torch.Tensor, - ss: torch.Tensor, - up: torch.Tensor, - hc_out: torch.Tensor, - norm_weight: torch.Tensor, - eps: float, - hc_count: int, - with_inject: bool, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - return launch(part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject) - - def reference( - self, - part: torch.Tensor, - ss: torch.Tensor, - up: torch.Tensor, - hc_out: torch.Tensor, - norm_weight: torch.Tensor, - eps: float, - hc_count: int, - with_inject: bool, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - """float64, one rounding per output.""" - dtype = hc_out.dtype - rank = up.shape[0] - hidden = hc_out.shape[1] // hc_count - rstd = torch.rsqrt(ss.double().t() / hidden + eps) - logits = torch.einsum("ts,stn->tn", rstd, part.double()) - low = F.silu(logits[:, :rank] / hc_count) - gate = torch.sigmoid(low @ up.double()) - x = hc_out.double().unflatten(-1, (hc_count, hidden)) * rstd.unsqueeze(-1) - normed = x.flatten(-2) * (1.0 + norm_weight.double()) - mixed = (gate * normed).unflatten(-1, (hc_count, hidden)).mean(dim=-2).to(dtype) - inject_out = None - if with_inject: - inject_out = (2.0 * torch.sigmoid(logits[:, rank : rank + hc_count] / hc_count)).to( - dtype - ) - return mixed, inject_out - - def is_valid( - self, - part: torch.Tensor, - ss: torch.Tensor, - up: torch.Tensor, - hc_out: torch.Tensor, - norm_weight: torch.Tensor, - eps: float, - hc_count: int, - with_inject: bool, - ) -> None: - rows, width = hc_out.shape - rank = up.shape[0] - assert hc_out.stride(1) == 1, "hc_out unit column stride" - assert part.dtype == torch.float32 and part.stride(2) == 1, "part fp32, unit last stride" - assert part.shape[:2] == (hc_count, rows), "part [C, T, N]" - assert part.shape[2] >= rank + (hc_count if with_inject else 0), "part too narrow" - assert ss.dtype == torch.float32 and ss.shape == (hc_count, rows), "ss fp32 [C, T]" - assert rows <= 1 or ss.stride(1) == 1, "ss unit stride along rows" - assert up.shape == (rank, width), "up [R, width]" - assert norm_weight.shape == (width,) and norm_weight.is_contiguous() - r2 = _rank_split(rank)[1] - assert r2 & (r2 - 1) == 0, "rank must be a power of two or a sum of two" - - -qwen4_exp_hc_decode_mix = Qwen4ExpHcDecodeMix() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix_fp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix_fp4.py deleted file mode 100644 index 4bb07863796a..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix_fp4.py +++ /dev/null @@ -1,431 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Decode-row hyper-connection site at an MoE site, second half: -`qwen4_exp_hc_decode_mix` (grouped RMS, SiLU bottleneck, injection weights, up -projection, sigmoid gate, stream mean) plus the NVFP4 quantization of the mean -for the routed experts -- one Triton GEMV launch. - -Tier 5 (Triton written for this target). Same structure as -`qwen4_exp_hc_decode_mix`, except that each program owns a whole number of -16-column NVFP4 blocks so it can quantize its own columns: the scale block is -the only cross-column dependency the quantization adds. - -The quantization helpers here (`nvfp4_quantize_blocks` in Triton, -`nvfp4_reference` in torch) reproduce `torch.ops.trtllm.fp4_quantize` (NVFP4, -linear scales) step for step, including its two `rcp.approx.ftz.f32` -reciprocals, so the packed data and the scales are bitwise that op's output on -the rounded mix; `qwen4_exp_hc_gate_mean_fp4` imports them. -""" - -from __future__ import annotations - -from typing import Optional, Tuple - -import torch -import torch.nn.functional as F -import triton -import triton.language as tl - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -from .qwen4_exp_hc_decode_mix import _rank_split, pdl_enabled - -#: Largest row count this entry is certified for. -MAX_ROWS = 16 - -#: NVFP4 block: one e4m3 scale per this many contiguous elements. -SF_VEC = 16 -_E4M3_MAX = 448.0 -_E2M1_MAX = 6.0 -_E2M1_MIDPOINTS = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0) -_TIE_UP = (False, True, False, True, False, True, False) -#: Same window `catalog/quantization/fp4_quantize` states for the same two -#: approximate reciprocals. -TIE_WINDOW = 2.0**-20 - - -@triton.jit -def rcp_approx_ftz(x): - """`rcp.approx.ftz.f32`, the reciprocal the trtllm quantizer uses.""" - return tl.inline_asm_elementwise( - "rcp.approx.ftz.f32 $0, $1;", "=r,r", [x], dtype=tl.float32, is_pure=True, pack=1 - ) - - -@triton.jit -def nvfp4_quantize_blocks(v, gs): - """NVFP4 of `v` `[NB, 8, 2]` fp32 (values already rounded to the activation - dtype; element `16*b + 2*p + j` of the row chunk at `[b, p, j]`) with the - global scale `gs` `[NB]`. Returns the packed bytes `[NB, 8]` (int32, element - `2k` in the low nibble) and the e4m3 scale bytes `[NB]` (uint8).""" - amax = tl.max(tl.max(tl.abs(v), axis=2), axis=1) - sfv = gs * (amax * rcp_approx_ftz(tl.zeros_like(amax) + 6.0)) - sf8 = sfv.to(tl.float8e4nv) - osc = rcp_approx_ftz(sf8.to(tl.float32) * rcp_approx_ftz(gs)) - osc = tl.where(amax != 0.0, osc, 0.0) - s = v * osc[:, None, None] - a = tl.abs(s) - # Round to nearest, ties to the even e2m1 code; saturate at 6. - code = ( - (a > 0.25).to(tl.int32) - + (a >= 0.75).to(tl.int32) - + (a > 1.25).to(tl.int32) - + (a >= 1.75).to(tl.int32) - + (a > 2.5).to(tl.int32) - + (a >= 3.5).to(tl.int32) - + (a > 5.0).to(tl.int32) - ) - code = code | tl.where(s.to(tl.int32, bitcast=True) < 0, 8, 0) - shift = 4 * tl.arange(0, 2) - packed = tl.sum(code << shift[None, None, :], axis=2) - return packed, sf8.to(tl.uint8, bitcast=True) - - -def nvfp4_reference( - x: torch.Tensor, global_scale: torch.Tensor -) -> Tuple[torch.Tensor, torch.Tensor]: - """Native-torch NVFP4 of `x` `[T, K]` with linear scales: per 16-block, - `sf = e4m3(gs * amax / 6)` and the data `x * gs / sf` rounded to e2m1 - (nearest, ties to even). Returns `(packed [T, K/2] uint8, sf [T*K/16] uint8)`.""" - packed, sf, _ = _nvfp4_torch(x, global_scale) - return packed, sf - - -def nvfp4_near_tie(x: torch.Tensor, global_scale: torch.Tensor) -> torch.Tensor: - """`[T, K]` bool: elements whose pre-rounding value sits within `TIE_WINDOW` - of an e2m1 midpoint, where approximate reciprocals may pick either side.""" - scaled = _nvfp4_torch(x, global_scale)[2] - mid = torch.tensor(_E2M1_MIDPOINTS, device=x.device, dtype=torch.float32) - a = scaled.abs().unsqueeze(-1) - return ((a - mid).abs() <= TIE_WINDOW * a).any(-1) - - -def _nvfp4_torch(x, global_scale): - gs = global_scale.float().reshape(()) - rows, k = x.shape - blocks = x.float().reshape(rows, k // SF_VEC, SF_VEC) - vecmax = blocks.abs().amax(-1) - sf = (gs * (vecmax / _E2M1_MAX)).clamp(max=_E4M3_MAX).to(torch.float8_e4m3fn) - out_scale = torch.where(vecmax != 0, gs / sf.float(), torch.zeros_like(vecmax)) - scaled = (blocks * out_scale.unsqueeze(-1)).reshape(rows, k) - mid = torch.tensor(_E2M1_MIDPOINTS, device=x.device, dtype=torch.float32) - tie_up = torch.tensor(_TIE_UP, device=x.device, dtype=torch.bool) - a = scaled.abs().unsqueeze(-1) - code = ((a > mid).sum(-1) + ((a == mid) & tie_up).any(-1)).to(torch.uint8) - negative = (scaled < 0) | ((scaled == 0) & torch.signbit(scaled)) - code = code | (negative.to(torch.uint8) << 3) - packed = code[:, 0::2] | (code[:, 1::2] << 4) - return packed, sf.view(torch.uint8).reshape(-1), scaled - - -@triton.jit -def _hc_decode_mix_fp4_kernel( - part_ptr, - ss_ptr, - up_ptr, - x_ptr, - norm_w_ptr, - gs_ptr, - mixed_ptr, - inj_out_ptr, - q_ptr, - sf_ptr, - stride_ps, - stride_pm, - stride_ss, - stride_up_r, - stride_up_n, - stride_x, - stride_mixed, - eps, - M: tl.constexpr, - HIDDEN: tl.constexpr, - HC: tl.constexpr, - R1: tl.constexpr, - R2: tl.constexpr, - BLOCK_H: tl.constexpr, - WITH_INJECT_OUT: tl.constexpr, - LAUNCH_WITH_PDL: tl.constexpr, - EARLY_TRIGGER: tl.constexpr, -): - NBH: tl.constexpr = BLOCK_H // 16 - pid = tl.program_id(0) - hh = pid * BLOCK_H + tl.arange(0, BLOCK_H) - sidx = tl.arange(0, HC) - col = sidx[None, :] * HIDDEN + hh[:, None] - hmask = (hh < HIDDEN)[:, None] & (sidx >= 0)[None, :] - ra = tl.arange(0, R1) - w_a = tl.load( - up_ptr + col[:, :, None] * stride_up_n + ra[None, None, :] * stride_up_r, - mask=hmask[:, :, None], - other=0.0, - ) - if R2 > 0: - rb = R1 + tl.arange(0, R2) - w_b = tl.load( - up_ptr + col[:, :, None] * stride_up_n + rb[None, None, :] * stride_up_r, - mask=hmask[:, :, None], - other=0.0, - ) - gam = 1.0 + tl.load(norm_w_ptr + col, mask=hmask, other=0.0).to(tl.float32) - b = pid * NBH + tl.arange(0, NBH) - p = tl.arange(0, 8) - bmask = b < HIDDEN // 16 - gs = tl.load(gs_ptr + tl.zeros([NBH], dtype=tl.int32)) - if LAUNCH_WITH_PDL: - if EARLY_TRIGGER: - tl.extra.cuda.gdc_launch_dependents() - tl.extra.cuda.gdc_wait() - for m in tl.static_range(M): - ssm = tl.load(ss_ptr + sidx * stride_ss + m) - pa = tl.load(part_ptr + sidx[:, None] * stride_ps + m * stride_pm + ra[None, :]) - if R2 > 0: - pb = tl.load(part_ptr + sidx[:, None] * stride_ps + m * stride_pm + rb[None, :]) - x = tl.load(x_ptr + m * stride_x + col, mask=hmask, other=0.0).to(tl.float32) - rstd = tl.rsqrt(ssm / HIDDEN + eps) - za = tl.sum(pa * rstd[:, None], axis=0) * (1.0 / HC) - acc = tl.sum(w_a.to(tl.float32) * (za * tl.sigmoid(za))[None, None, :], axis=2) - if R2 > 0: - zb = tl.sum(pb * rstd[:, None], axis=0) * (1.0 / HC) - acc += tl.sum(w_b.to(tl.float32) * (zb * tl.sigmoid(zb))[None, None, :], axis=2) - v = tl.sigmoid(acc) * (x * rstd[None, :] * gam) - mixed = (tl.sum(v, axis=1) * (1.0 / HC)).to(mixed_ptr.dtype.element_ty) - tl.store(mixed_ptr + m * stride_mixed + hh, mixed, mask=hh < HIDDEN) - packed, sf = nvfp4_quantize_blocks(tl.reshape(mixed.to(tl.float32), [NBH, 8, 2]), gs) - q_off = m * (HIDDEN // 2) + b[:, None] * 8 + p[None, :] - tl.store(q_ptr + q_off, packed.to(tl.uint8), mask=bmask[:, None]) - tl.store(sf_ptr + m * (HIDDEN // 16) + b, sf, mask=bmask) - if WITH_INJECT_OUT: - ci = tl.arange(0, HC) - pi = tl.load( - part_ptr + sidx[:, None] * stride_ps + m * stride_pm + (R1 + R2) + ci[None, :], - mask=(sidx >= 0)[:, None] & (pid == 0), - other=0.0, - ) - zi = tl.sum(pi * rstd[:, None], axis=0) * (1.0 / HC) - tl.store( - inj_out_ptr + m * HC + ci, - (2.0 * tl.sigmoid(zi)).to(inj_out_ptr.dtype.element_ty), - mask=(ci >= 0) & (pid == 0), - ) - - -#: Launch configuration: (BLOCK_H, num_warps). BLOCK_H a multiple of 16. Swept -#: on GB300 at 1/2/4 rows (graph, cold weights): (16, 4) fastest at every row -#: count; (16, 2), (16, 8) and (32, 4) 0.3..5 us slower. -DEFAULT_CONFIG = (16, 4) -#: Let the kernel after this one launch as soon as every program has requested its -#: weights: a programmatic dependent (the decode projections) then streams its own -#: weight while this kernel finishes. Its `gdc_wait` still waits for this whole grid. -TRIGGER_DEPENDENTS = True - - -def launch( - part, - ss, - up, - hc_out, - norm_weight, - eps, - hc_count, - with_inject, - global_scale, - config=DEFAULT_CONFIG, - trigger=TRIGGER_DEPENDENTS, -): - """The launch behind `Qwen4ExpHcDecodeMixFp4.__call__`, tile config exposed. - Returns `(mixed, inject_out, xq, xsf)`.""" - rows, width = hc_out.shape - hidden = width // hc_count - rank = up.shape[0] - r1, r2 = _rank_split(rank) - block_h, warps = config - dev, dt = hc_out.device, hc_out.dtype - mixed = torch.empty((rows, hidden), dtype=dt, device=dev) - inject_out = torch.empty((rows, hc_count), dtype=dt, device=dev) if with_inject else None - xq = torch.empty((rows, hidden // 2), dtype=torch.uint8, device=dev) - xsf = torch.empty((rows * (hidden // SF_VEC),), dtype=torch.uint8, device=dev) - pdl = pdl_enabled(dev.index if dev.index is not None else torch.cuda.current_device()) - if rows > 0: - _hc_decode_mix_fp4_kernel[(triton.cdiv(hidden, block_h),)]( - part, - ss, - up, - hc_out, - norm_weight, - global_scale, - mixed, - inject_out if with_inject else mixed, - xq, - xsf, - part.stride(0), - part.stride(1), - ss.stride(0), - up.stride(0), - up.stride(1), - hc_out.stride(0), - mixed.stride(0), - eps, - M=rows, - HIDDEN=hidden, - HC=hc_count, - R1=r1, - R2=r2, - BLOCK_H=block_h, - WITH_INJECT_OUT=with_inject, - LAUNCH_WITH_PDL=pdl, - EARLY_TRIGGER=trigger, - num_warps=warps, - launch_pdl=pdl, - ) - return mixed, inject_out, xq, xsf - - -class Qwen4ExpHcDecodeMixFp4(OpWrapper): - """`qwen4_exp_hc_decode_mix` plus the NVFP4 of its `mixed`. With `C = - hc_count`, `H = width // C`, `R = up.shape[0]`, `part` / `ss` the per-stream - partial logits `[C, T, >= R (+ C)]` and sums of squares `[C, T]` from - `qwen4_exp_hc_decode_first`, `up` the `[R, width]` view of the up weight:: - - rstd[t, s] = 1 / sqrt(ss[s, t] / H + eps) - logits[t, n] = sum_s rstd[t, s] * part[s, t, n] - low = silu(logits[:, :R] / C) - gate = sigmoid(low @ up) # [T, C*H] - normed[t, s*H+i] = hc_out[t, s*H+i] * rstd[t, s] * (1 + norm_weight[s*H+i]) - mixed[t, i] = mean over s of gate[t, s*H+i] * normed[t, s*H+i] - inject_out = 2 * sigmoid(logits[:, R:R+C] / C) # with_inject only - xq, xsf = NVFP4(mixed) # e2m1 pairs, e4m3 scale per 16, linear layout - - fp32 inside, one rounding per output; the quantization reads the rounded - `mixed`. Returns fresh `(mixed, inject_out, xq [T, H/2] uint8, xsf [T*H/16] - uint8)`; `inject_out` is None without `with_inject`. `global_scale` is the - 1-element fp32 NVFP4 global scale. - - Fusion boundary: everything `qwen4_exp_hc_decode_mix` fuses, plus the NVFP4 - quantization of `mixed`. The partials are the producer's. - """ - - ARCHS = frozenset({Arch.SM_103}) - - # spec: rows, with_inject (True), hc_count (4), hidden (2560), rank (320), dtype (bf16). - CELLS: tuple[Cell, ...] = tuple( - Cell(why=f"checkpoint 4x2560 r320 bf16, inject, {why}", spec=dict(rows=t)) - for t, why in ( - (1, "decode 1 row (con 1)"), - (2, "decode 2 rows"), - (3, "decode 3 rows"), - (16, "decode 16 rows (MAX_ROWS)"), - ) - ) + ( - Cell(why="hc_out a row slice buf[::2], 4 rows", spec=dict(rows=4, layout="row_slice")), - Cell(why="hc_count 2 at hidden 2560, 2 rows", spec=dict(rows=2, hc_count=2)), - Cell( - why="hidden 144 rank 32 (9 scale blocks), 5 rows", - spec=dict(rows=5, hidden=144, rank=32), - ), - Cell(why="no inject, 1 row", spec=dict(rows=1, with_inject=False)), - Cell(why="float16 operands, 4 rows", spec=dict(rows=4, dtype=torch.float16)), - Cell(why="zero rows", spec=dict(rows=0)), - ) - - note = """ - Metadata consumed: none. Call sites: `DecodeTarget`'s trunk, the MoE-site - mixer of every layer (48 per forward) at batches of at most - `_HC_DECODE_GEMV_MAX_ROWS` rows, right after `qwen4_exp_hc_decode_first`; - `xq`/`xsf` feed `fp4_block_scale_moe_renorm_routed` (linear scales). - - Preconditions: those of `qwen4_exp_hc_decode_mix`, plus `H` a multiple of - 16 and `global_scale` one fp32 element on the device. Row count is a - compile-time constant, at most MAX_ROWS certified. - - Numerics: `mixed` and `inject_out` follow `qwen4_exp_hc_decode_mix`'s - formula (fp32, one rounding); `inject_out` is bitwise its result, `mixed` - can differ in a few elements by one rounding because the 16-column program - tile reduces the up projection in another order. The NVFP4 is - the trtllm recipe (two `rcp.approx.ftz.f32`, `cvt.rn.satfinite` e4m3 - scale, e2m1 round-to-nearest-even saturating at 6) on the rounded `mixed`; - the GPU test checks it bitwise against `torch.ops.trtllm.fp4_quantize`. - Graph-capturable; programmatic dependent launch as `qwen4_exp_hc_decode_mix`. - - Grid: cdiv(H, BLOCK_H) programs (BLOCK_H a multiple of 16; 16 -> 160 - programs of 4 warps, each one [16, C, R] weight load, 40 KB). Measured on - GB300 (graph replay, cold weights, 1 row): 4.6 us, against 6.0 us for - `qwen4_exp_hc_decode_mix` + `fp4_quantize` and 3.9 us for - `qwen4_exp_hc_decode_mix` alone (its 8-column tile has twice the programs). - """ - - def __call__( - self, - part: torch.Tensor, - ss: torch.Tensor, - up: torch.Tensor, - hc_out: torch.Tensor, - norm_weight: torch.Tensor, - eps: float, - hc_count: int, - with_inject: bool, - global_scale: torch.Tensor, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, torch.Tensor]: - return launch(part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject, global_scale) - - def reference( - self, - part: torch.Tensor, - ss: torch.Tensor, - up: torch.Tensor, - hc_out: torch.Tensor, - norm_weight: torch.Tensor, - eps: float, - hc_count: int, - with_inject: bool, - global_scale: torch.Tensor, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, torch.Tensor]: - """float64 with one rounding per output, then the NVFP4 of the rounded - `mixed` in native torch.""" - dtype = hc_out.dtype - rank = up.shape[0] - hidden = hc_out.shape[1] // hc_count - rstd = torch.rsqrt(ss.double().t() / hidden + eps) - logits = torch.einsum("ts,stn->tn", rstd, part.double()) - low = F.silu(logits[:, :rank] / hc_count) - gate = torch.sigmoid(low @ up.double()) - x = hc_out.double().unflatten(-1, (hc_count, hidden)) * rstd.unsqueeze(-1) - normed = x.flatten(-2) * (1.0 + norm_weight.double()) - mixed = (gate * normed).unflatten(-1, (hc_count, hidden)).mean(dim=-2).to(dtype) - inject_out = None - if with_inject: - inject_out = (2.0 * torch.sigmoid(logits[:, rank : rank + hc_count] / hc_count)).to( - dtype - ) - return (mixed, inject_out) + nvfp4_reference(mixed, global_scale) - - def is_valid( - self, - part: torch.Tensor, - ss: torch.Tensor, - up: torch.Tensor, - hc_out: torch.Tensor, - norm_weight: torch.Tensor, - eps: float, - hc_count: int, - with_inject: bool, - global_scale: torch.Tensor, - ) -> None: - rows, width = hc_out.shape - rank = up.shape[0] - assert hc_out.stride(1) == 1, "hc_out unit column stride" - assert part.dtype == torch.float32 and part.stride(2) == 1, "part fp32, unit last stride" - assert part.shape[:2] == (hc_count, rows), "part [C, T, N]" - assert part.shape[2] >= rank + (hc_count if with_inject else 0), "part too narrow" - assert ss.dtype == torch.float32 and ss.shape == (hc_count, rows), "ss fp32 [C, T]" - assert rows <= 1 or ss.stride(1) == 1, "ss unit stride along rows" - assert up.shape == (rank, width), "up [R, width]" - assert norm_weight.shape == (width,) and norm_weight.is_contiguous() - r2 = _rank_split(rank)[1] - assert r2 & (r2 - 1) == 0, "rank must be a power of two or a sum of two" - assert (width // hc_count) % SF_VEC == 0, "hidden must be whole 16-blocks" - # The kernel loads element 0 only; a per-row scale would be applied as row 0's. - assert global_scale.numel() == 1 and global_scale.dtype == torch.float32 - - -qwen4_exp_hc_decode_mix_fp4 = Qwen4ExpHcDecodeMixFp4() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix_fp4_cute.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix_fp4_cute.py index a12ca131fe3b..30cd06f39156 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix_fp4_cute.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix_fp4_cute.py @@ -5,12 +5,9 @@ partials, the up projection, the sigmoid gate, the stream mean and the NVFP4 of the mean for the routed experts -- one PDL launch. -Tier 5 (CuTe DSL written for this target). `qwen4_exp_hc_decode_mix_fp4` computes the -same for up to 16 rows in Triton, whose programs hold a 64-row x 320 tile of the up -weight in 128 threads. Here a CTA -owns one 16-column NVFP4 block of the hidden width across the 4 streams (64 rows of -the up weight), and 8 threads share a row, 40 weights each; 160 CTAs of 512 threads, -all resident at once. +Tier 5 (CuTe DSL written for this target). A CTA owns one 16-column NVFP4 block of the +hidden width across the 4 streams (64 rows of the up weight), and 8 threads share a +row, 40 weights each; 160 CTAs of 512 threads, all resident at once. Order of work inside a CTA: @@ -33,15 +30,16 @@ import cutlass import cutlass.cute as cute import torch +import torch.nn.functional as F from cutlass import const_expr from cutlass.cute import experimental as cute_ext from cutlass.cute.runtime import from_dlpack from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -from .qwen4_exp_hc_decode_first_cute import pdl_enabled -from .qwen4_exp_hc_decode_mix import qwen4_exp_hc_decode_mix as _triton_mix -from .qwen4_exp_hc_decode_mix_fp4 import qwen4_exp_hc_decode_mix_fp4 as _triton_mix_fp4 +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( + fp4_quantize, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl MAX_ROWS = 1 VEC = 8 @@ -294,7 +292,7 @@ def launch(part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject, global xsf = torch.empty((rows * (hidden // 16),), dtype=torch.uint8, device=dev) if fp4 else None if rows == 0: return mixed, inject_out, xq, xsf - pdl = pdl_enabled(dev.index if dev.index is not None else torch.cuda.current_device()) + pdl = get_env_enable_pdl() part2 = part[:, 0, :] fn = _compiled( hidden, @@ -334,16 +332,42 @@ def launch(part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject, global return mixed, inject_out, xq, xsf -def _check(part, ss, up, hc_out, norm_weight, hc_count) -> None: +def _check(part, ss, up, hc_out, norm_weight, hc_count, with_inject) -> None: rows, width = hc_out.shape + rank = up.shape[0] + assert hc_out.stride(1) == 1, "hc_out unit column stride" + assert part.dtype == torch.float32 and part.stride(2) == 1, "part fp32, unit last stride" + assert part.shape[:2] == (hc_count, rows), "part [C, T, N]" + assert part.shape[2] >= rank + (hc_count if with_inject else 0), "part too narrow" + assert ss.dtype == torch.float32 and ss.shape == (hc_count, rows), "ss fp32 [C, T]" + assert up.shape == (rank, width), "up [R, width]" + assert norm_weight.shape == (width,) and norm_weight.is_contiguous() assert rows <= MAX_ROWS, "one row" - assert width // hc_count == 2560 and up.shape[0] == 320 and hc_count == 4, ( + assert width // hc_count == 2560 and rank == 320 and hc_count == 4, ( "the checkpoint's 4 x 2560, rank 320" ) assert hc_out.dtype == torch.bfloat16 and up.dtype == torch.bfloat16, "bf16" assert up.t().is_contiguous(), "up the [R, width] view of a row-major [width, R] weight" +def _reference(part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject): + """`(mixed, inject_out)` in float64, one rounding per output.""" + dtype = hc_out.dtype + rank = up.shape[0] + hidden = hc_out.shape[1] // hc_count + rstd = torch.rsqrt(ss.double().t() / hidden + eps) + logits = torch.einsum("ts,stn->tn", rstd, part.double()) + low = F.silu(logits[:, :rank] / hc_count) + gate = torch.sigmoid(low @ up.double()) + x = hc_out.double().unflatten(-1, (hc_count, hidden)) * rstd.unsqueeze(-1) + normed = x.flatten(-2) * (1.0 + norm_weight.double()) + mixed = (gate * normed).unflatten(-1, (hc_count, hidden)).mean(dim=-2).to(dtype) + inject_out = None + if with_inject: + inject_out = (2.0 * torch.sigmoid(logits[:, rank : rank + hc_count] / hc_count)).to(dtype) + return mixed, inject_out + + _CELLS = ( Cell(why="checkpoint 4x2560 r320 bf16, inject, 1 row (con 1)", spec=dict(rows=1)), Cell(why="stack-final mixer (no inject), 1 row", spec=dict(rows=1, with_inject=False)), @@ -352,10 +376,27 @@ def _check(part, ss, up, hc_out, norm_weight, hc_count) -> None: class Qwen4ExpHcDecodeMixFp4Cute(OpWrapper): - """`qwen4_exp_hc_decode_mix_fp4` for a one-row batch: `(mixed, inject_out, xq, - xsf)`, the NVFP4 of the rounded `mixed` with linear scales (see that entry). - - Fusion boundary: as `qwen4_exp_hc_decode_mix_fp4`. + """`qwen4_exp_hc_decode_mix_cute` plus the NVFP4 of its `mixed` for the routed + experts. With `C = hc_count`, `H = width // C`, `R = up.shape[0]`, `part` / `ss` the + per-stream partial logits `[C, T, >= R (+ C)]` and sums of squares `[C, T]` from + `qwen4_exp_hc_decode_first_cute`, `up` the `[R, width]` view of the up weight:: + + rstd[t, s] = 1 / sqrt(ss[s, t] / H + eps) + logits[t, n] = sum_s rstd[t, s] * part[s, t, n] + low = silu(logits[:, :R] / C) + gate = sigmoid(low @ up) # [T, C*H] + normed[t, s*H+i] = hc_out[t, s*H+i] * rstd[t, s] * (1 + norm_weight[s*H+i]) + mixed[t, i] = mean over s of gate[t, s*H+i] * normed[t, s*H+i] + inject_out = 2 * sigmoid(logits[:, R:R+C] / C) # with_inject only + xq, xsf = NVFP4(mixed) # e2m1 pairs, e4m3 scale per 16, linear layout + + fp32 inside, one rounding per output; the quantization reads the rounded `mixed`. + Returns fresh `(mixed, inject_out, xq [T, H/2] uint8, xsf [T*H/16] uint8)`; + `inject_out` is None without `with_inject`. `global_scale` is the 1-element fp32 + NVFP4 global scale. + + Fusion boundary: everything `qwen4_exp_hc_decode_mix_cute` fuses, plus the NVFP4 + quantization of `mixed`. The partials are the producer's. """ ARCHS = frozenset({Arch.SM_103}) @@ -363,35 +404,40 @@ class Qwen4ExpHcDecodeMixFp4Cute(OpWrapper): note = """ Metadata consumed: none. Call sites: `DecodeTarget`'s MoE sites at a one-row - decode step, after `qwen4_exp_hc_decode_first_cute`; batches of more rows keep the - Triton entries. - - Preconditions: the Triton entry's, and one row, the checkpoint's 4 x 2560 width, - rank 320, bf16; `up` the `.t()` view of a contiguous `[width, R]` weight; - `global_scale` fp32 with one element. - - Numerics: fp32 inside, one rounding per output, the Triton entry's formula and - order up to fp32 reassociation of the up-projection dot (inside its band); the - NVFP4 is `fp4_quantize`'s arithmetic on the rounded `mixed` (`rcp.approx.ftz` - reciprocals, e4m3 scale rounded to nearest, e2m1 nearest-even with saturation at - 6). Deterministic; graph-capturable once compiled (the first call per variant). + decode step, after `qwen4_exp_hc_decode_first_cute`; `xq`/`xsf` feed + `fp4_block_scale_moe_renorm_routed` (linear scales). Batches of more rows take + `qwen4_exp_hc_up_gate_mean_fp4`. + + Preconditions: those of `qwen4_exp_hc_decode_mix_cute`, and `global_scale` fp32 with + one element. + + Numerics: `mixed` and `inject_out` as `qwen4_exp_hc_decode_mix_cute`; the NVFP4 is + `fp4_quantize`'s arithmetic on the rounded `mixed` (`rcp.approx.ftz` reciprocals, + e4m3 scale rounded to nearest, e2m1 nearest-even with saturation at 6). + Deterministic; graph-capturable once compiled (the first call per variant). PDL: the weight rows and norm elements stream before `griddepcontrol.wait`. """ def __call__(self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject, global_scale): return launch(part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject, global_scale) - def reference(self, *args, **kwargs): - """The Triton entry's reference: float64, then the NVFP4 of the rounded `mixed`.""" - return _triton_mix_fp4.reference(*args, **kwargs) + def reference( + self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject, global_scale + ): + """float64 with one rounding per output, then `fp4_quantize`'s native-torch NVFP4 + of the rounded `mixed` (linear scales).""" + mixed, inject_out = _reference( + part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject + ) + return (mixed, inject_out) + fp4_quantize.reference( + mixed, global_scale, BLOCK_H, False, False + ) def is_valid( self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject, global_scale ) -> None: - _triton_mix_fp4.is_valid( - part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject, global_scale - ) - _check(part, ss, up, hc_out, norm_weight, hc_count) + _check(part, ss, up, hc_out, norm_weight, hc_count, with_inject) + # The kernel loads element 0 only; a per-row scale would be applied as row 0's. assert global_scale.numel() == 1 and global_scale.dtype == torch.float32, ( "global_scale fp32 [1]" ) @@ -401,10 +447,16 @@ def is_valid( class Qwen4ExpHcDecodeMixCute(OpWrapper): - """`qwen4_exp_hc_decode_mix` for a one-row batch: `(mixed, inject_out)` -- this - kernel's plain build (no NVFP4 output). - - Fusion boundary: as `qwen4_exp_hc_decode_mix`. + """The second half of a one-row decode hyper-connection site: the grouped RMS, the + SiLU bottleneck and the injection weights from `qwen4_exp_hc_decode_first_cute`'s + per-stream partials, then the up projection, the sigmoid gate on the normed stream + and the mean over streams -- `qwen4_exp_hc_decode_mix_fp4_cute`'s formulas without + the NVFP4 output. Returns fresh `(mixed, inject_out)` in `hc_out`'s dtype; + `inject_out` is None without `with_inject`. + + Fusion boundary: the grouped RMS, the bottleneck SiLU, the injection weights, the + up projection, the sigmoid gate, the normed stream and the stream mean. The + partials are the producer's. """ ARCHS = frozenset({Arch.SM_103}) @@ -412,17 +464,19 @@ class Qwen4ExpHcDecodeMixCute(OpWrapper): note = """ Metadata consumed: none. Call sites: `DecodeTarget`'s non-MoE sites at a one-row - decode step, after `qwen4_exp_hc_decode_first_cute`; batches of more rows keep the - Triton entry. Beside a first kernel capped at 96 registers, its 160 CTAs (512 - threads at 49 registers) are all resident while that kernel runs, so the up weight - is in registers when the wait returns; the Triton entry's 320 programs are not. - - Preconditions: the Triton entry's, and one row, the checkpoint's 4 x 2560 width, - rank 320, bf16; `up` the `.t()` view of a contiguous `[width, R]` weight. - - Numerics: fp32 inside, one rounding per output, the Triton entry's formula and - order up to fp32 reassociation of the up-projection dot (inside its band). - Deterministic; graph-capturable once compiled (the first call per variant). + decode step, after `qwen4_exp_hc_decode_first_cute`; batches of more rows take + `qwen4_exp_hc_up_gate_mean`. Beside a first kernel capped at 96 registers, its 160 + CTAs (512 threads at 49 registers) are all resident while that kernel runs, so the + up weight is in registers when the wait returns. + + Preconditions: `part` fp32 `[C, T, >= R (+ C)]` and `ss` fp32 `[C, T]`, `part` with + unit last stride; `hc_out` `[T, width]` with unit column stride; `norm_weight` + contiguous `[width]`; one row, the checkpoint's 4 x 2560 width, rank 320, bf16; + `up` the `.t()` view of a contiguous `[width, R]` weight. + + Numerics: fp32 inside (`low`, the gate logits and `normed` are never rounded), one + rounding per output, inside the float64 reference's band. Deterministic; + graph-capturable once compiled (the first call per variant). """ def __call__(self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject): @@ -431,13 +485,12 @@ def __call__(self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject ) return mixed, inject_out - def reference(self, *args, **kwargs): - """The Triton entry's float64 reference.""" - return _triton_mix.reference(*args, **kwargs) + def reference(self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject): + """float64, one rounding per output.""" + return _reference(part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject) def is_valid(self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject) -> None: - _triton_mix.is_valid(part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject) - _check(part, ss, up, hc_out, norm_weight, hc_count) + _check(part, ss, up, hc_out, norm_weight, hc_count, with_inject) qwen4_exp_hc_decode_mix_cute = Qwen4ExpHcDecodeMixCute() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_tc_first.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_tc_first.py index e476bd967313..dfed495b80ca 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_tc_first.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_tc_first.py @@ -23,8 +23,6 @@ from __future__ import annotations -import os - import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute @@ -44,13 +42,14 @@ PipelineUmmaAsync, ) from tensorrt_llm._torch.cute_dsl_kernels.blackwell.utils import make_ptr +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl #: Largest row count this entry is certified for: one 128-row MMA tile. MAX_ROWS = 128 #: Programmatic dependent launch (sm_90+), off with TRTLLM_ENABLE_PDL=0. Without it #: the grid-dependency wait and the dependents trigger are no-ops. -_PDL = os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" +_PDL = get_env_enable_pdl() def _asm(ret, args, text, constraints, loc, ip, side_effects=False): @@ -597,7 +596,7 @@ class Qwen4ExpHcDecodeTcFirst(OpWrapper): that `qwen4_exp_hc_lowrank_act` reads back (three launches), and the Triton `tl.dot` K-loop designs measured 36-60 us at 128 rows. Metadata consumed: none. Call sites: `DecodeTarget`'s trunk - hyper-connection mixers from `_HC_TC_FIRST_MIN_ROWS` to + hyper-connection mixers past `_HC_DECODE_GEMV_MAX_ROWS` and up to `_HC_TC_FIRST_MAX_ROWS` (= MAX_ROWS) rows, 97 per forward, between `qwen4_exp_hc_inject_norm` (or the PLE tail) and `qwen4_exp_hc_up_gate_mean(_fp4)`. Preconditions (bf16): `normed` @@ -614,10 +613,10 @@ class Qwen4ExpHcDecodeTcFirst(OpWrapper): before `griddepcontrol.wait`, `normed` after it; dependents released at kernel start. Graph-capturable (current stream, no host sync, no persistent state). Compiles once per (BN, KS, C, R, N, K, inject). - Measured on GB300 (graph replay, 24 cold weights, `microbench/` of the - candidate): 6.4 us at 128 rows against 8.6 us for cuBLAS split-K + - reduce + lowrank_act; tile sweep BN 32 / 48 / 64 / 112 x KS 8 / 10 / - 16, 48 x 16 fastest at every row count. + Measured on GB300 (graph replay, 24 cold weights): 6.4 us at 128 rows + against 8.6 us for cuBLAS split-K + reduce + lowrank_act; tile sweep + BN 32 / 48 / 64 / 112 x KS 8 / 10 / 16, 48 x 16 fastest at every row + count. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_gate_mean.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_gate_mean.py deleted file mode 100644 index 0aaa1bd5bcd5..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_gate_mean.py +++ /dev/null @@ -1,125 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Hyper-connection mixer epilogue: the sigmoid gate applied to the normed -stream and the mean over streams, from the up projection's logits, in one -Triton kernel. - -Tier 5 (Triton written for this target). Used where the up projection runs on -cuBLAS (decode below `_HC_FUSED_UP_MIN_ROWS` rows): `qwen4_exp_hc_lowrank_act` -> -`cublas_mm(low, up^T)` -> this. Elsewhere `qwen4_exp_hc_up_gate_mean` computes the -up projection and this epilogue in one kernel. -""" - -from __future__ import annotations - -import torch -import triton -import triton.language as tl - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - - -@triton.jit -def _hc_gate_mean_kernel( - gate_ptr, - normed_ptr, - mixed_ptr, - stride_gate, - stride_normed, - stride_mixed, - HIDDEN: tl.constexpr, - HC: tl.constexpr, - BLOCK: tl.constexpr, -): - row = tl.program_id(0).to(tl.int64) - h = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) - mask = h < HIDDEN - acc = tl.zeros([BLOCK], dtype=tl.float32) - for s in tl.static_range(HC): - g = tl.load(gate_ptr + row * stride_gate + s * HIDDEN + h, mask=mask, other=0.0) - x = tl.load(normed_ptr + row * stride_normed + s * HIDDEN + h, mask=mask, other=0.0) - acc += tl.sigmoid(g.to(tl.float32)) * x.to(tl.float32) - tl.store( - mixed_ptr + row * stride_mixed + h, - (acc * (1.0 / HC)).to(mixed_ptr.dtype.element_ty), - mask=mask, - ) - - -class Qwen4ExpHcGateMean(OpWrapper): - """With `C = hc_count`, `H = width // C`, `gate_logits` and `normed` both - `[T, C*H]` stream-major:: - - mixed[t, i] = mean over s of sigmoid(gate_logits[t, s*H + i]) * normed[t, s*H + i] - - fp32 inside, one rounding. Returns a fresh `[T, H]` tensor in `normed`'s - dtype. Fusion boundary: the sigmoid, the product and the stream mean; the - up projection producing `gate_logits` is the caller's (`cublas_mm`). - """ - - ARCHS = frozenset({Arch.SM_103}) - - # spec: rows, hc_count (4), hidden (2560), dtype (bf16). - CELLS: tuple[Cell, ...] = tuple( - Cell(why=f"checkpoint 4x2560 bf16, {why}", spec=dict(rows=t)) - for t, why in ( - (1, "1 row"), - (128, "128 rows"), - (2048, "prefill 2048"), - (8192, "prefill 8192"), - ) - ) + ( - Cell(why="hc_count 2, 64 rows", spec=dict(rows=64, hc_count=2)), - Cell(why="hc_count 8, 64 rows", spec=dict(rows=64, hc_count=8)), - Cell(why="hidden 128 (not the checkpoint), 48 rows", spec=dict(rows=48, hidden=128)), - Cell(why="float16 operands, 64 rows", spec=dict(rows=64, dtype=torch.float16)), - Cell(why="zero rows", spec=dict(rows=0)), - ) - - note = """ - Metadata consumed: none. Call sites: the trunk's hyper-connection sites in - `DecodeTarget` below `_HC_FUSED_UP_MIN_ROWS` rows, after `cublas_mm(low, up^T)`. Preconditions: both - operands 2-D `[T, C*H]` with unit column stride. Graph-capturable. - Measured (GB300, graph): the three-launch cuBLAS route this ends is 2.7x - faster than `qwen4_exp_hc_gate_mix` at 512..8192 rows (106 vs 287 us at - 8192) and 1.7x at 128. - """ - - def __call__( - self, gate_logits: torch.Tensor, normed: torch.Tensor, hc_count: int - ) -> torch.Tensor: - rows, width = normed.shape - hidden = width // hc_count - mixed = torch.empty((rows, hidden), dtype=normed.dtype, device=normed.device) - if rows > 0: - block = min(1024, triton.next_power_of_2(hidden)) - _hc_gate_mean_kernel[(rows, triton.cdiv(hidden, block))]( - gate_logits, - normed, - mixed, - gate_logits.stride(0), - normed.stride(0), - mixed.stride(0), - HIDDEN=hidden, - HC=hc_count, - BLOCK=block, - num_warps=4, - ) - return mixed - - def reference( - self, gate_logits: torch.Tensor, normed: torch.Tensor, hc_count: int - ) -> torch.Tensor: - """float64, one rounding.""" - hidden = normed.shape[1] // hc_count - gate = torch.sigmoid(gate_logits.double()).unflatten(-1, (hc_count, hidden)) - x = normed.double().unflatten(-1, (hc_count, hidden)) - return (gate * x).mean(dim=-2).to(normed.dtype) - - def is_valid(self, gate_logits: torch.Tensor, normed: torch.Tensor, hc_count: int) -> None: - assert gate_logits.shape == normed.shape, "gate_logits and normed must match" - assert gate_logits.stride(1) == 1 and normed.stride(1) == 1, "unit column stride" - assert normed.shape[1] % hc_count == 0 - - -qwen4_exp_hc_gate_mean = Qwen4ExpHcGateMean() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_inject_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_inject_norm.py index 347942e58938..5b3b3cf75b28 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_inject_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_inject_norm.py @@ -22,7 +22,6 @@ from __future__ import annotations -import os import struct from typing import Optional, Tuple @@ -37,10 +36,11 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper from tensorrt_llm._torch.cute_dsl_kernels.blackwell.utils import make_ptr +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl #: Trigger dependents at kernel start (`griddepcontrol.launch_dependents`); a no-op #: unless the next kernel is launched with programmatic dependent launch. -_PDL_TRIGGER = os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" +_PDL_TRIGGER = get_env_enable_pdl() #: Rows from which a wide (`hidden >= 2048`) launch runs two warps a program #: instead of eight: prefill-sized grids are bandwidth-bound and run faster diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_lowrank_act.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_lowrank_act.py index ba61c79dad97..5c2b72a06ded 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_lowrank_act.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_lowrank_act.py @@ -4,9 +4,10 @@ the rank-`R` bottleneck and the block-injection weights, from the first projection's fp32 logits, in one Triton row kernel. -Tier 5 (Triton written for this target). Used where the up projection runs on -cuBLAS (prefill-sized row counts), between `cublas_mm` (first projection) and -`cublas_mm` (up projection) + `qwen4_exp_hc_gate_mean`. +Tier 5 (Triton written for this target). Used where the first projection runs on +cuBLAS (prefill, and decode batches past `qwen4_exp_hc_decode_tc_first`'s rows), +between `cublas_mm` (first projection) and the fused up projection +`qwen4_exp_hc_up_gate_mean` (or its NVFP4 variant). """ from __future__ import annotations diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean.py index 1ee0d815777f..4e0e9a5eded2 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean.py @@ -17,8 +17,6 @@ from __future__ import annotations -import os - import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute @@ -38,9 +36,10 @@ PipelineUmmaAsync, ) from tensorrt_llm._torch.cute_dsl_kernels.blackwell.utils import make_ptr +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl #: Programmatic dependent launch (sm_90+), off with TRTLLM_ENABLE_PDL=0. -_PDL = os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" +_PDL = get_env_enable_pdl() @dsl_user_op @@ -829,8 +828,8 @@ class Qwen4ExpHcUpGateMean(OpWrapper): The GEMM accumulates in fp32 in tensor memory and is never rounded or written; sigmoid, product and stream mean are fp32 with one rounding to `normed`'s dtype. Returns a fresh `[T, H]` tensor. Fusion boundary: the up - projection, the sigmoid gate and the stream mean (what `cublas_mm(low, up)` - + `qwen4_exp_hc_gate_mean` compute as two launches). + projection, the sigmoid gate and the stream mean (two launches as `cublas_mm(low, + up)` followed by a separate gate-and-mean pass). """ ARCHS = frozenset({Arch.SM_103}) @@ -859,8 +858,7 @@ class Qwen4ExpHcUpGateMean(OpWrapper): Tier 6 (CuTe DSL written for this target): a persistent tcgen05 GEMM whose epilogue applies the gate, so the `[T, C*H]` gate logits never reach HBM; tiers 1-5 have no GEMM with a custom multi-input epilogue (cuBLAS writes - the logits; the Triton `qwen4_exp_hc_gate_mix` was 2.7x slower than - cuBLAS + gate_mean at prefill rows). Metadata consumed: none. Call sites: + the logits). Metadata consumed: none. Call sites: the trunk's hyper-connection mixers (`_hc_mixer`). Preconditions (bf16): `low` `[T, R]` contiguous, `up` the transposed view of a contiguous `[C*H, R]` weight, `normed` `[T, C*H]` contiguous, 16-byte aligned diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean_fp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean_fp4.py index 654d5263d1c3..601bfef7d350 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean_fp4.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean_fp4.py @@ -20,10 +20,15 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( + fp4_quantize, +) -from .qwen4_exp_hc_decode_mix_fp4 import SF_VEC, nvfp4_reference from .qwen4_exp_hc_up_gate_mean import _BLOCK_H, hc_up_gate_mean_fp4, qwen4_exp_hc_up_gate_mean +#: NVFP4 block: one e4m3 scale per this many contiguous elements. +SF_VEC = 16 + class Qwen4ExpHcUpGateMeanFp4(OpWrapper): """With `C = hc_count`, `H = width // C`, `low` the `[T, R]` SiLU bottleneck, @@ -69,7 +74,7 @@ class Qwen4ExpHcUpGateMeanFp4(OpWrapper): kernel compiled with its NVFP4 epilogue; tiers 1-5 have no GEMM with this multi-input epilogue, and the fused up projection itself is tier 6. Metadata consumed: none. Call sites: the 48 MoE-site mixers (prefill, and - decode from `_HC_FUSED_UP_MIN_ROWS` rows); `xq`/`xsf` feed + decode past `_HC_DECODE_GEMV_MAX_ROWS` rows); `xq`/`xsf` feed `fp4_block_scale_moe_renorm_routed` (linear scales), `mixed` the router and the shared expert. Preconditions: those of `qwen4_exp_hc_up_gate_mean` (bf16, `R % 64 == 0`, `H % 32 == 0`, `C * 32 <= 256`, 16-byte aligned @@ -104,7 +109,7 @@ def reference( """`mixed` in float64 with one rounding, then the NVFP4 of that rounded value in native torch.""" mixed = qwen4_exp_hc_up_gate_mean.reference(low, up, normed, hc_count) - return (mixed,) + nvfp4_reference(mixed, global_scale) + return (mixed,) + fp4_quantize.reference(mixed, global_scale, SF_VEC, False, False) def is_valid( self, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_grouped_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_grouped_gemm.py index 81216c78ecf0..ed325ce10d52 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_grouped_gemm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_grouped_gemm.py @@ -12,8 +12,7 @@ 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 -#: e2m1 magnitudes by 3-bit code. -_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) +from .cute_dsl_nvfp4_gather_grouped_gemm_act_fusion import _dequant, unswizzle def _ckpt( @@ -430,28 +429,4 @@ def is_valid( ) -# ── native-torch NVFP4 helpers for reference() and the test's compare ──────── - - -def unswizzle(buf: torch.Tensor, rows: int, cols: int) -> torch.Tensor: - """128x4-swizzled scale bytes -> `[rows, cols]` (rows % 128 == 0, cols % 4 == 0). - - Byte `(r, c)` sits at `(r // 128) * 128 * cols + (c // 4) * 512 + (r % 32) * 16 - + ((r % 128) // 32) * 4 + c % 4`. - """ - return buf.reshape(rows // 128, cols // 4, 32, 4, 4).permute(0, 3, 2, 1, 4).reshape(rows, cols) - - -def _dequant(packed: torch.Tensor, sf: torch.Tensor, vec: int) -> torch.Tensor: - """`[R, K/2]` packed e2m1 (element 2j in the low nibble) and `[R, K/vec]` e4m3 - bytes -> float64 `[R, K]`. Exact.""" - nibble = torch.tensor(_E2M1, dtype=torch.float64, device=packed.device) - nibble = torch.cat([nibble, -nibble]) - byte = torch.arange(256, device=packed.device) - pair = torch.stack([nibble[byte & 0xF], nibble[byte >> 4]], dim=1) - vals = pair[packed.long()].reshape(packed.shape[0], -1) - scales = sf.to(torch.uint8).view(torch.float8_e4m3fn).double() - return vals * scales.repeat_interleave(vec, dim=1) - - cute_dsl_nvfp4_grouped_gemm = CuteDslNvfp4GroupedGemm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_gated_norm_pdl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_gated_norm_pdl.py index 1686f08195a9..02d928a6a24c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_gated_norm_pdl.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_gated_norm_pdl.py @@ -3,8 +3,8 @@ """One Gated-DeltaNet decode step per row, and the gated RMS norm on its output, in one PDL launch. -Tier 5 (CuTe DSL written for this target; the recurrence is the gated delta rule -as `gated_delta_rule_decode_pretranspose` states it, reimplemented here). At one +Tier 5 (CuTe DSL written for this target; the recurrence is the gated delta rule of +FlashInfer's `gdn_decode.gated_delta_rule_decode_pretranspose`, reimplemented here). At one decode row the step reads and writes each v-head's 64 KB fp32 state once and otherwise does a few hundred flops per state row; the rest of its time is latency. The state is not produced by anything in the step ahead of this kernel -- the @@ -24,9 +24,6 @@ from __future__ import annotations -import functools -import os - import cuda.bindings.driver as cuda_drv import cutlass import cutlass.cute as cute @@ -38,9 +35,7 @@ from cutlass.cutlass_dsl import dsl_user_op from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.gated_delta_rule_decode_pretranspose import ( # noqa: E501 - gated_delta_rule_decode_pretranspose as _recurrence, -) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl from ..norm.rms_norm_gated_token_major import rms_norm_gated_token_major as _gated_norm @@ -51,13 +46,6 @@ WARPS = 4 -@functools.lru_cache(maxsize=None) -def pdl_enabled(device_index: int) -> bool: - """Programmatic dependent launch: sm_90+ unless TRTLLM_ENABLE_PDL=0.""" - props = torch.cuda.get_device_properties(device_index) - return props.major >= 9 and os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" - - @dsl_user_op def _st_shared_cluster_f32(ptr, val, *, loc=None, ip=None): """`st.shared::cluster.f32` to a pointer from `cute.arch.map_dsmem_ptr`.""" @@ -295,29 +283,19 @@ def _dl(t: torch.Tensor, align: int = 4): return from_dlpack(t, assumed_align=align) -@functools.lru_cache(maxsize=None) -def _compiled(rows, k_heads, v_heads, kdim, vdim, proj_stride, pool_shape, pool_stride, pdl): - kern = _GdnKernel(k_heads=k_heads, v_heads=v_heads, kdim=kdim, vdim=vdim, pdl=pdl) - dev = "cuda" - width = 2 * k_heads * kdim + 2 * v_heads * vdim + 2 * v_heads - proj = torch.empty(rows, proj_stride, dtype=torch.bfloat16, device=dev)[:, :width] - pool = torch.empty_strided(pool_shape, pool_stride, dtype=torch.float32, device=dev) - slot = torch.empty(rows, dtype=torch.int32, device=dev) - a_log_h = torch.empty(v_heads, dtype=torch.float32, device=dev) - dt = torch.empty(v_heads, dtype=torch.float32, device=dev) - w = torch.empty(vdim, dtype=torch.bfloat16, device=dev) - out = torch.empty(rows, v_heads * vdim, dtype=torch.bfloat16, device=dev) - stream = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream) - return cute_ext.compile( - kern, +_COMPILED: dict = {} + + +def _operands(proj, pool, slots, A_log, dt_bias, norm_weight, out, eps, stream): + return ( _dl(proj, 8), _dl(pool, 16), - _dl(slot), - _dl(a_log_h), - _dl(dt), - _dl(w, 2), + _dl(slots), + _dl(A_log), + _dl(dt_bias), + _dl(norm_weight, 2), _dl(out, 2), - cutlass.Float32(1e-6), + cutlass.Float32(eps), stream, ) @@ -330,9 +308,10 @@ def launch(proj, pool, slots, A_log, dt_bias, norm_weight, eps, k_heads, kdim, p out = torch.empty((rows, v_heads * vdim), dtype=torch.bfloat16, device=proj.device) if rows == 0: return out - dev = proj.device - pdl = pdl and pdl_enabled(dev.index if dev.index is not None else torch.cuda.current_device()) - fn = _compiled( + pdl = bool(pdl and get_env_enable_pdl()) + stream = cuda_drv.CUstream(torch.cuda.current_stream(proj.device).cuda_stream) + operands = (proj, pool, slots, A_log, dt_bias, norm_weight, out, eps, stream) + key = ( rows, k_heads, v_heads, @@ -340,24 +319,53 @@ def launch(proj, pool, slots, A_log, dt_bias, norm_weight, eps, k_heads, kdim, p vdim, proj.stride(0), tuple(pool.shape), - tuple(pool.stride()), - bool(pdl), - ) - stream = cuda_drv.CUstream(torch.cuda.current_stream(dev).cuda_stream) - fn( - _dl(proj, 8), - _dl(pool, 16), - _dl(slots), - _dl(A_log), - _dl(dt_bias), - _dl(norm_weight, 2), - _dl(out, 2), - cutlass.Float32(eps), - stream, + pool.stride(), + pdl, ) + fn = _COMPILED.get(key) + if fn is None: + # Compiled from the call's own operands (a stand-in for the pool would span every + # layer of the coalesced state cache). + kern = _GdnKernel(k_heads=k_heads, v_heads=v_heads, kdim=kdim, vdim=vdim, pdl=pdl) + fn = _COMPILED[key] = cute_ext.compile(kern, *_operands(*operands)) + fn(*_operands(*operands)) return out +def _recurrence_reference(q, k, v, A_log, a, dt_bias, b, pool, slots): + """`(o, pool_after)` of one gated-delta-rule decode step per row in fp64, rounded once: + `o` to bf16, the state to the pool's dtype. Row `n` reads and writes `pool[slots[n]]` + (a negative slot skips the row: zero output, no state access); q and k are + L2-normalized, the decay is `exp(-exp(A_log) * softplus(a + dt_bias))`, `beta = + sigmoid(b)`, the scale `K^-0.5`. `q` / `k` `[n, 1, H, K]`, `v` `[n, 1, HV, V]`, `a` / + `b` `[n, 1, HV]`, `pool` `[slots, HV, V, K]`. Mutates nothing.""" + dt = torch.float64 + batch, _, num_k_heads, k_dim = q.shape + hv = v.shape[2] + rep = hv // num_k_heads + scale = k_dim**-0.5 + reads = slots.tolist() + pool_after = pool.clone() + o = torch.zeros(batch, 1, hv, v.shape[3], dtype=dt, device=q.device) + neg_a = -torch.exp(A_log.to(dt)) + for n in range(batch): + if reads[n] < 0: + continue + q_n = q[n, 0].to(dt).repeat_interleave(rep, dim=0) + k_n = k[n, 0].to(dt).repeat_interleave(rep, dim=0) + q_n = q_n * torch.rsqrt(q_n.pow(2).sum(-1, keepdim=True) + 1e-6) + k_n = k_n * torch.rsqrt(k_n.pow(2).sum(-1, keepdim=True) + 1e-6) + x = a[n, 0].to(dt) + dt_bias.to(dt) + softplus = torch.where(x <= 20.0, torch.log1p(torch.exp(x)), x) + beta = torch.sigmoid(b[n, 0].to(dt)) + h = pool[reads[n]].to(dt) * torch.exp(neg_a * softplus)[:, None, None] + u = (v[n, 0].to(dt) - torch.einsum("hvk,hk->hv", h, k_n)) * beta[:, None] + h = h + u[:, :, None] * k_n[:, None, :] + o[n, 0] = torch.einsum("hvk,hk->hv", h, q_n * scale) + pool_after[reads[n]] = h.to(pool.dtype) + return o.to(torch.bfloat16), pool_after + + class GdnDecodeGatedNormPdl(OpWrapper): """A single-token Gated-DeltaNet step per row over the fp32 slot pool, then the sigmoid-gated RMS norm of its output per v-head. With `proj` the decode input @@ -406,9 +414,8 @@ class GdnDecodeGatedNormPdl(OpWrapper): Launch: `HV * 8` CTAs a row, clusters of 8 (one v-head), 128 threads, <= 64 registers. Numerics: fp32 inside; `o` rounded to bf16 before the norm, the output - rounded once. Inside the recurrence entry's band (torch's default for bf16 / fp32) of the - fp64 composition. Deterministic. Graph-capturable after the first call per - geometry (compiled then). + rounded once. Inside torch's default band for bf16 / fp32 of the fp64 composition. + Deterministic. Graph-capturable after the first call per geometry (compiled then). """ def __call__( @@ -439,7 +446,7 @@ def reference( kdim: int, pdl: bool = True, ) -> tuple[torch.Tensor, torch.Tensor]: - """`(out, pool_after)`: the recurrence entry's fp64 reference (o rounded to + """`(out, pool_after)`: the fp64 recurrence (`_recurrence_reference`; o rounded to bf16, the pool to fp32), then the gated norm entry's fp32 reference.""" rows = proj.shape[0] v_heads, vdim = pool.shape[1], pool.shape[2] @@ -448,21 +455,16 @@ def reference( [k_heads * kdim, k_heads * kdim, v_heads * vdim, v_heads * vdim, v_heads, v_heads], 1, ) - o, pool_after = _recurrence.reference( - q=q.reshape(rows, 1, k_heads, kdim), - k=k.reshape(rows, 1, k_heads, kdim), - v=v.reshape(rows, 1, v_heads, vdim), - state=None, - A_log=A_log, - a=a.reshape(rows, 1, v_heads), - dt_bias=dt_bias, - b=b.reshape(rows, 1, v_heads), - scale=None, - output=torch.empty(rows, 1, v_heads, vdim, dtype=torch.bfloat16, device=proj.device), - use_qk_l2norm=True, - initial_state=pool, - initial_state_indices=slots, - output_state_indices=None, + o, pool_after = _recurrence_reference( + q.reshape(rows, 1, k_heads, kdim), + k.reshape(rows, 1, k_heads, kdim), + v.reshape(rows, 1, v_heads, vdim), + A_log, + a.reshape(rows, 1, v_heads), + dt_bias, + b.reshape(rows, 1, v_heads), + pool, + slots, ) out = _gated_norm.reference( o.reshape(-1, vdim), z.reshape(rows, v_heads, vdim), norm_weight, eps, None, "sigmoid" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_rows_gated_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_rows_gated_norm.py index 33c82416a042..b81e768bb14a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_rows_gated_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_rows_gated_norm.py @@ -3,10 +3,10 @@ """One Gated-DeltaNet decode step per row and the gated RMS norm on its output, streamed over the fp32 state pool: the decode step at batch sizes past the cluster kernel's. -Tier 6 (CuTe DSL written for this target; the recurrence is the gated delta rule as -`gated_delta_rule_decode_pretranspose` states it, reimplemented here). At a decode batch -the step is the pool's traffic -- each (row, v-head) state, 64 KB fp32, read and written -once -- so a CTA owns one whole state and streams it: warp w the value rows +Tier 6 (CuTe DSL written for this target; the recurrence is the gated delta rule of +FlashInfer's `gdn_decode.gated_delta_rule_decode_pretranspose`, reimplemented here). At a +decode batch the step is the pool's traffic -- each (row, v-head) state, 64 KB fp32, read +and written once -- so a CTA owns one whole state and streams it: warp w the value rows [w * RPW, (w + 1) * RPW), lane l the key columns [4l, 4l + 4) of each, so a warp's access is one contiguous 512-byte state row. The state is loaded before `griddepcontrol.wait` (nothing in the step writes it ahead of this kernel). The two dot @@ -29,13 +29,11 @@ from cutlass.cutlass_dsl import T, dsl_user_op from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.gated_delta_rule_decode_pretranspose import ( # noqa: E501 - gated_delta_rule_decode_pretranspose as _recurrence, -) from tensorrt_llm._torch.cute_dsl_kernels.blackwell.utils import make_ptr +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl -from ..norm.rms_norm_gated_token_major import rms_norm_gated_token_major as _gated_norm -from .gdn_decode_gated_norm_pdl import pdl_enabled +from .gdn_decode_gated_norm_pdl import _recurrence_reference +from .gdn_decode_gated_norm_pdl import gdn_decode_gated_norm_pdl as _one_launch #: Key columns a lane (a 16-byte vector of a state row). _PL = 4 @@ -388,10 +386,7 @@ def launch(proj, pool, slots, A_log, dt_bias, norm_weight, eps, k_heads, kdim, p out = torch.empty((rows, v_heads * vdim), dtype=torch.bfloat16, device=proj.device) if rows == 0: return out - dev = proj.device - pdl = bool( - pdl and pdl_enabled(dev.index if dev.index is not None else torch.cuda.current_device()) - ) + pdl = bool(pdl and get_env_enable_pdl()) key = (k_heads, v_heads, kdim, vdim, pdl, norm) ptrs = ( proj, @@ -497,8 +492,8 @@ class GdnDecodeRowsGatedNorm(OpWrapper): Numerics: fp32 inside; each dot product over K is a lane's 4-term FMA chain and then the transposed reduction's tree; `o` rounded to bf16 before the norm, the output - rounded once. Inside the recurrence entry's band (torch's default for bf16 / fp32) - of the composition's float64 reference, not bitwise FlashInfer's kernel. + rounded once. Inside torch's default band for bf16 / fp32 of the composition's + float64 reference, not bitwise FlashInfer's kernel. """ def __call__( @@ -529,35 +524,12 @@ def reference( kdim: int, pdl: bool = True, ) -> tuple[torch.Tensor, torch.Tensor]: - """`(out, pool_after)`: the recurrence entry's fp64 reference (o rounded to - bf16, the pool to fp32), then the gated norm entry's fp32 reference.""" - rows = proj.shape[0] - v_heads, vdim = pool.shape[1], pool.shape[2] - q, k, v, z, a, b = torch.split( - proj, - [k_heads * kdim, k_heads * kdim, v_heads * vdim, v_heads * vdim, v_heads, v_heads], - 1, + """`(out, pool_after)`: `gdn_decode_gated_norm_pdl`'s reference, the same contract -- + the fp64 recurrence (o rounded to bf16, the pool to fp32), then the gated norm + entry's fp32 reference.""" + return _one_launch.reference( + proj, pool, slots, A_log, dt_bias, norm_weight, eps, k_heads, kdim, pdl ) - o, pool_after = _recurrence.reference( - q=q.reshape(rows, 1, k_heads, kdim), - k=k.reshape(rows, 1, k_heads, kdim), - v=v.reshape(rows, 1, v_heads, vdim), - state=None, - A_log=A_log, - a=a.reshape(rows, 1, v_heads), - dt_bias=dt_bias, - b=b.reshape(rows, 1, v_heads), - scale=None, - output=torch.empty(rows, 1, v_heads, vdim, dtype=torch.bfloat16, device=proj.device), - use_qk_l2norm=True, - initial_state=pool, - initial_state_indices=slots, - output_state_indices=None, - ) - out = _gated_norm.reference( - o.reshape(-1, vdim), z.reshape(rows, v_heads, vdim), norm_weight, eps, None, "sigmoid" - ) - return out.reshape(rows, v_heads * vdim), pool_after def is_valid( self, @@ -655,7 +627,7 @@ def reference( out: torch.Tensor, pdl: bool = True, ) -> tuple[torch.Tensor, torch.Tensor]: - """`(o, pool_after)`: the recurrence entry's fp64 reference, o rounded to bf16.""" + """`(o, pool_after)`: the fp64 recurrence, o rounded to bf16 (the pool to fp32).""" rows = proj.shape[0] v_heads, vdim = pool.shape[1], pool.shape[2] q, k, v, _, a, b = torch.split( @@ -663,21 +635,16 @@ def reference( [k_heads * kdim, k_heads * kdim, v_heads * vdim, v_heads * vdim, v_heads, v_heads], 1, ) - o, pool_after = _recurrence.reference( - q=q.reshape(rows, 1, k_heads, kdim), - k=k.reshape(rows, 1, k_heads, kdim), - v=v.reshape(rows, 1, v_heads, vdim), - state=None, - A_log=A_log, - a=a.reshape(rows, 1, v_heads), - dt_bias=dt_bias, - b=b.reshape(rows, 1, v_heads), - scale=None, - output=torch.empty(rows, 1, v_heads, vdim, dtype=torch.bfloat16, device=proj.device), - use_qk_l2norm=True, - initial_state=pool, - initial_state_indices=slots, - output_state_indices=None, + o, pool_after = _recurrence_reference( + q.reshape(rows, 1, k_heads, kdim), + k.reshape(rows, 1, k_heads, kdim), + v.reshape(rows, 1, v_heads, vdim), + A_log, + a.reshape(rows, 1, v_heads), + dt_bias, + b.reshape(rows, 1, v_heads), + pool, + slots, ) return o.reshape(rows, v_heads * vdim), pool_after diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_in_proj_conv1d_pdl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_in_proj_conv1d_pdl.py index de6fd934aaa1..b31ddad3daab 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_in_proj_conv1d_pdl.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_in_proj_conv1d_pdl.py @@ -29,9 +29,6 @@ from __future__ import annotations -import functools -import os - import cuda.bindings.driver as cuda_drv import cutlass import cutlass.cute as cute @@ -48,6 +45,7 @@ OpWrapper, true_fp32_matmul, ) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl VEC = 8 #: Before its PDL wait a CTA also has the weight rows of CTAs `bidx + j * L2_DIST` @@ -84,13 +82,6 @@ def _bulk_prefetch_l2(src, nbytes, *, loc=None, ip=None): _COMPILE_OPTIONS = "--ptxas-options -maxrregcount=64" -@functools.lru_cache(maxsize=None) -def pdl_enabled(device_index: int) -> bool: - """Programmatic dependent launch: sm_90+ unless TRTLLM_ENABLE_PDL=0.""" - props = torch.cuda.get_device_properties(device_index) - return props.major >= 9 and os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" - - class _InProjConvKernel: def __init__(self, *, rows, k, threads, outs, conv_dim, pad_slot, pdl): assert k % (threads * VEC) == 0 and threads % 32 == 0 @@ -245,30 +236,18 @@ def _dl(t: torch.Tensor, align: int): return from_dlpack(t, assumed_align=align) -@functools.lru_cache(maxsize=None) -def _compiled(rows, n, k, conv_dim, pad_slot, pdl, state_shape, state_stride): - threads, outs = TILING[k] - kern = _InProjConvKernel( - rows=rows, k=k, threads=threads, outs=outs, conv_dim=conv_dim, pad_slot=pad_slot, pdl=pdl - ) - dev = "cuda" - a = torch.empty(rows, k, dtype=torch.bfloat16, device=dev) - w = torch.empty(n, k, dtype=torch.bfloat16, device=dev) - out = torch.empty(rows, n, dtype=torch.bfloat16, device=dev) - state = torch.empty_strided(state_shape, state_stride, dtype=torch.bfloat16, device=dev) - convw = torch.empty(conv_dim, WIDTH, dtype=torch.bfloat16, device=dev) - slot = torch.empty(rows, dtype=torch.int32, device=dev) - stream = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream) - return cute_ext.compile( - kern, - _dl(a, 32), - _dl(w, 32), +_COMPILED: dict = {} + + +def _operands(x, mat_b, out, conv_state, conv_weight, conv_state_indices, stream): + return ( + _dl(x, 32), + _dl(mat_b.t(), 32), _dl(out, 32), - _dl(state, 2), - _dl(convw, 8), - _dl(slot, 4), + _dl(conv_state, 2), + _dl(conv_weight, 8), + _dl(conv_state_indices, 4), stream, - options=_COMPILE_OPTIONS, ) @@ -277,28 +256,38 @@ def launch(x, mat_b, conv_state, conv_weight, conv_state_indices, conv_dim, pad_ rows, k = x.shape n = mat_b.shape[1] dev = x.device - pdl = pdl and pdl_enabled(dev.index if dev.index is not None else torch.cuda.current_device()) + pdl = bool(pdl and get_env_enable_pdl()) out = torch.empty(rows, n, dtype=x.dtype, device=dev) - fn = _compiled( + stream = cuda_drv.CUstream(torch.cuda.current_stream(dev).cuda_stream) + operands = (x, mat_b, out, conv_state, conv_weight, conv_state_indices, stream) + key = ( rows, n, k, conv_dim, int(pad_slot_id), - bool(pdl), + pdl, tuple(conv_state.shape), - tuple(conv_state.stride()), - ) - stream = cuda_drv.CUstream(torch.cuda.current_stream(dev).cuda_stream) - fn( - _dl(x, 32), - _dl(mat_b.t(), 32), - _dl(out, 32), - _dl(conv_state, 2), - _dl(conv_weight, 8), - _dl(conv_state_indices, 4), - stream, + conv_state.stride(), ) + fn = _COMPILED.get(key) + if fn is None: + # Compiled from the call's own operands (a stand-in for the pool would span every + # layer of the coalesced conv cache). + threads, outs = TILING[k] + kern = _InProjConvKernel( + rows=rows, + k=k, + threads=threads, + outs=outs, + conv_dim=conv_dim, + pad_slot=int(pad_slot_id), + pdl=pdl, + ) + fn = _COMPILED[key] = cute_ext.compile( + kern, *_operands(*operands), options=_COMPILE_OPTIONS + ) + fn(*_operands(*operands)) return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_prefill_conv_post.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_prefill_conv_post.py index 06b7863fc652..fe94663b51c0 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_prefill_conv_post.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_prefill_conv_post.py @@ -6,7 +6,7 @@ in, and each sequence's conv window published to the paged pool. Tier 6 (CuTe DSL written for this target) for the conv pass, Tier 5 (Triton) -for the gates, the window publish and the fallback conv pass. The three-call +for the gates and the window publish. The three-call composition it replaces -- `extract_transpose_prefill_slice` (token-major -> channel-major copy), `causal_conv1d_fwd` (channel-major conv, in place) and `fused_gdn_post_conv` (split + norm + gates) -- moves the `[T, 10240]` @@ -43,142 +43,10 @@ DT_BIAS_RANGE = (-8.9375, 19.25) SOFTPLUS_THRESHOLD = 20.0 -#: Launch shape of the main kernel, chosen by micro-bench at 2048, 8192 and -#: 16384 tokens (see the entry note). The register cap is what makes it fast: -#: uncapped, the compiler takes 140+ registers for the pool-reading branch and -#: the occupancy that leaves cannot hide the load latency; capped at 72 that -#: branch spills, and it runs only on blocks holding a sequence's first rows. -_BLOCK_T = 32 -_NUM_WARPS = 4 -_MAXNREG = 72 #: Channels per program of the window publish. _PUBLISH_BLOCK_C = 512 -@triton.jit(do_not_specialize=["num_tokens", "num_seqs", "search_steps"]) -def _gdn_prefill_conv_post_kernel( - x_ptr, - stride_xt, - w_ptr, - stride_wc, - stride_wk, - state_ptr, - stride_ss, - stride_sc, - stride_sw, - qsl_ptr, - slot_ptr, - init_ptr, - a_ptr, - stride_at, - stride_ah, - b_ptr, - stride_bt, - stride_bh, - A_log_ptr, - dt_bias_ptr, - q_ptr, - k_ptr, - v_ptr, - g_ptr, - beta_ptr, - num_tokens, - num_seqs, - search_steps, - pad_slot_id, - l2_norm_eps, - softplus_threshold, - IDX: tl.constexpr, - NUM_K_HEADS: tl.constexpr, - NUM_V_HEADS: tl.constexpr, - HEAD_DIM: tl.constexpr, - BLOCK_T: tl.constexpr, -): - """One program: BLOCK_T tokens x one 128-channel head of the conv output.""" - head = tl.program_id(0) - t = tl.program_id(1) * BLOCK_T + tl.arange(0, BLOCK_T) - t_mask = t < num_tokens - d = tl.arange(0, HEAD_DIM) - c = head * HEAD_DIM + d - rows2 = (tl.program_id(1) * BLOCK_T + tl.arange(0, BLOCK_T))[:, None] - rmask2 = rows2 < num_tokens - - # Sequence of each token: the largest s with qsl[s] <= t (a zero-length - # sequence shares its start with the next one and is never chosen over it). - lo = tl.zeros([BLOCK_T], dtype=tl.int32) - hi = tl.zeros([BLOCK_T], dtype=tl.int32) + num_seqs - for _ in range(search_steps): - mid = (lo + hi) // 2 - right = tl.load(qsl_ptr + mid) <= t - lo = tl.where(right, mid, lo) - hi = tl.where(right, hi, mid) - s0 = tl.load(qsl_ptr + lo) - - # Width-4 causal conv, taps accumulated in fp32 in tap order. A tap before - # the sequence start reads the carried window column W-1-(s0-pos), or 0. - # Only a block holding one of a sequence's first three tokens needs that; - # every other block takes the branch whose taps are plain row loads. - acc = tl.zeros([BLOCK_T, HEAD_DIM], dtype=tl.float32) - if tl.min(tl.where(t_mask, t - s0, 3), axis=0) < 3: - slot = tl.load(slot_ptr + lo, mask=t_mask, other=pad_slot_id) - carry = t_mask & (tl.load(init_ptr + lo, mask=t_mask, other=0) != 0) & (slot != pad_slot_id) - slot64 = slot.to(tl.int64) - for kk in tl.static_range(4): - pos = t - 3 + kk - in_seq = pos >= s0 - xv = tl.load( - x_ptr + pos[:, None].to(IDX) * stride_xt + c[None, :], - mask=(t_mask & in_seq)[:, None], - other=0.0, - ) - s_row = slot64 * stride_ss + (3 + pos - s0) * stride_sw - sv = tl.load( - state_ptr + s_row[:, None] + c[None, :] * stride_sc, - mask=(carry & (~in_seq))[:, None], - other=0.0, - ) - wk = tl.load(w_ptr + c * stride_wc + kk * stride_wk).to(tl.float32) - acc += tl.where(in_seq[:, None], xv, sv).to(tl.float32) * wk[None, :] - else: - for kk in tl.static_range(4): - xv = tl.load( - x_ptr + (rows2 - (3 - kk)).to(IDX) * stride_xt + c[None, :], mask=rmask2, other=0.0 - ) - wk = tl.load(w_ptr + c * stride_wc + kk * stride_wk).to(tl.float32) - acc += xv.to(tl.float32) * wk[None, :] - act = acc / (1.0 + tl.exp(-acc)) - # The conv result is rounded to the activation dtype once, as the - # unfused conv stores it, before anything reads it back. - y = act.to(q_ptr.dtype.element_ty) - - ri2 = rows2.to(IDX) - if head < 2 * NUM_K_HEADS: - yf = y.to(tl.float32) - inv = 1.0 / tl.sqrt(tl.sum(yf * yf, axis=1) + l2_norm_eps) - normed = (yf * inv[:, None]).to(q_ptr.dtype.element_ty) - o_off = ri2 * (NUM_K_HEADS * HEAD_DIM) + (head % NUM_K_HEADS) * HEAD_DIM + d[None, :] - if head < NUM_K_HEADS: - tl.store(q_ptr + o_off, normed, mask=rmask2) - else: - tl.store(k_ptr + o_off, normed, mask=rmask2) - else: - vh = head - 2 * NUM_K_HEADS - o_off = ri2 * (NUM_V_HEADS * HEAD_DIM) + vh * HEAD_DIM + d[None, :] - tl.store(v_ptr + o_off, y, mask=rmask2) - t64 = t.to(IDX) - a_val = tl.load(a_ptr + t64 * stride_at + vh * stride_ah, mask=t_mask, other=0.0) - b_val = tl.load(b_ptr + t64 * stride_bt + vh * stride_bh, mask=t_mask, other=0.0) - A_log = tl.load(A_log_ptr + vh).to(tl.float32) - dt_bias = tl.load(dt_bias_ptr + vh).to(tl.float32) - gate_input = a_val.to(tl.float32) + dt_bias - softplus = tl.where( - gate_input <= softplus_threshold, tl.log(1.0 + tl.exp(gate_input)), gate_input - ) - g_off = t64 * NUM_V_HEADS + vh - tl.store(g_ptr + g_off, -tl.exp(A_log) * softplus, mask=t_mask) - tl.store(beta_ptr + g_off, tl.sigmoid(b_val.to(tl.float32)), mask=t_mask) - - @triton.jit def _gdn_prefill_conv_publish_kernel( x_ptr, @@ -233,7 +101,7 @@ def _gdn_prefill_conv_publish_kernel( # registers; then a thread (4 channels x `_TPT` consecutive tokens) slides the width-4 # window over shared-memory rows, every element converted to fp32 once. The conv runs # as `fma.rn.f32x2` on channel pairs (first tap from zero, then in tap order: per element -# the Triton pass's `fma.rn.f32` chain); the SiLU is `acc * rcp(1 + 2^(-acc * log2e))` +# an `fma.rn.f32` chain); the SiLU is `acc * rcp(1 + 2^(-acc * log2e))` # with `ex2.approx.ftz` and `rcp.approx.ftz`, rounded in pairs by `cvt.rn.bf16x2.f32`. # -------------------------------------------------------------------------------------- @@ -307,9 +175,9 @@ def _fmul2(a0, a1, b0, b1, *, loc=None, ip=None): @dsl_user_op def _silu2(a0, a1, *, loc=None, ip=None): """`acc * rcp(1 + 2^(acc * -log2e))` for two values: `mul.rn.f32x2`, `ex2.approx.ftz`, - `add.rn.f32x2`, `rcp.approx.ftz`, `mul.rn`. For `1 + e` in [1, 2^126] this equals the - Triton pass's `div.full.f32`; past it (acc < -87.3) it gives -0 where that gives a - value below 1e-36.""" + `add.rn.f32x2`, `rcp.approx.ftz`, `mul.rn`. For `1 + e` in [1, 2^126] this equals a + Triton `acc / (1 + exp(-acc))` (`div.full.f32`); past it (acc < -87.3) it gives -0 + where that gives a value below 1e-36.""" r = _asm2( T.f32(), [cutlass.Float32(a0), cutlass.Float32(a1)], @@ -343,7 +211,8 @@ def _fma_rn(a, b, c, *, loc=None, ip=None): @dsl_user_op def _inv_norm(ss, eps, *, loc=None, ip=None): - """1 / sqrt(ss + eps): `add.f32`, `sqrt.approx.ftz.f32`, `div.full.f32` (the Triton pass's).""" + """1 / sqrt(ss + eps): `add.f32`, `sqrt.approx.ftz.f32`, `div.full.f32` (a Triton + `1.0 / tl.sqrt(...)`).""" return cutlass.Float32( _asm1( T.f32(), @@ -686,8 +555,7 @@ def _gdn_gates_kernel( HP: tl.constexpr, BLOCK_T: tl.constexpr, ): - """The decay/beta gates over `[T, NUM_V_HEADS]`, one row-coalesced tile a program: - the Triton conv pass's code for them.""" + """The decay/beta gates over `[T, NUM_V_HEADS]`, one row-coalesced tile a program.""" t = (tl.program_id(0) * BLOCK_T + tl.arange(0, BLOCK_T))[:, None] h = tl.arange(0, HP)[None, :] m = (t < num_tokens) & (h < NUM_V_HEADS) @@ -802,9 +670,6 @@ def _launch( pad_slot_id, l2_norm_eps, softplus_threshold, - block_t=_BLOCK_T, - num_warps=_NUM_WARPS, - maxnreg=_MAXNREG, ): assert head_k_dim == head_v_dim and weight.shape[-1] == CONV_WIDTH, ( "one program per 128-channel head: head_k_dim must equal head_v_dim, and the " @@ -821,9 +686,12 @@ def _launch( beta = torch.empty((1, num_tokens, num_v_heads), dtype=torch.float32, device=dev) if num_seqs == 0: return q, k, v, g, beta - if num_tokens > 0 and _cute_conv_pass_ok( - x, weight, conv_states, num_k_heads, head_k_dim, num_v_heads - ): + if num_tokens > 0: + assert _cute_conv_pass_ok(x, weight, conv_states, num_k_heads, head_k_dim, num_v_heads), ( + "the conv pass takes bf16 x with 128-channel heads, 16-byte aligned rows and " + f"int32 weight/pool strides; got {x.dtype}, head_k_dim {head_k_dim}, x at " + f"{x.data_ptr() % 16} mod 16 with row stride {x.stride(0)}" + ) _cute_conv_pass( x, weight, conv_states, query_start_loc, cache_indices, has_initial_state, q, k, v, num_k_heads, num_v_heads, head_k_dim, pad_slot_id, l2_norm_eps, @@ -846,50 +714,6 @@ def _launch( BLOCK_T=64, num_warps=4, ) - elif num_tokens > 0: - heads = 2 * num_k_heads + num_v_heads - blocks = triton.cdiv(num_tokens, block_t) - # int32 element offsets suffice while the largest one fits. - out_width = max(num_k_heads * head_k_dim, num_v_heads * head_v_dim) - off64 = (num_tokens + 3) * max(x.stride(0), a.stride(0), b.stride(0), out_width) >= 2**31 - _gdn_prefill_conv_post_kernel[(heads, blocks)]( - x, - x.stride(0), - weight, - weight.stride(0), - weight.stride(1), - conv_states, - *conv_states.stride(), - query_start_loc, - cache_indices, - has_initial_state, - a, - a.stride(0), - a.stride(1), - b, - b.stride(0), - b.stride(1), - A_log, - dt_bias, - q, - k, - v, - g, - beta, - num_tokens, - num_seqs, - max(num_seqs - 1, 0).bit_length() + 1, - pad_slot_id, - l2_norm_eps, - softplus_threshold, - IDX=tl.int64 if off64 else tl.int32, - NUM_K_HEADS=num_k_heads, - NUM_V_HEADS=num_v_heads, - HEAD_DIM=head_k_dim, - BLOCK_T=block_t, - num_warps=num_warps, - maxnreg=maxnreg, - ) _gdn_prefill_conv_publish_kernel[(num_seqs, triton.cdiv(conv_dim, _PUBLISH_BLOCK_C))]( x, x.stride(0), @@ -1076,29 +900,21 @@ class GdnPrefillConvPost(OpWrapper): inside one sequence takes a straight-line branch, the rest read the pool token by token. The gates (Triton): one row-coalesced `[64, 48]` tile a program. Then a per-sequence window publish (Triton), after both, so no - program reads a window another has already advanced. Inputs the CuTe - pass cannot take (not bf16, a row of `x` not 16-byte aligned, other head - sizes) run the Triton conv pass (`_gdn_prefill_conv_post_kernel`, gates - included) instead. + program reads a window another has already advanced. Speed (one layer, 461-slot pool, in_proj views, us, graph replay over 4 - layers' inputs): 16384 tokens / 3 contexts 136 (the Triton conv pass 198, - the three-call chain it replaced 528). The Triton pass is bound by its - instruction stream (ncu: 183 M warp instructions, 74% issue-busy, DRAM - 41%); the CuTe pass's straight-line window and paired fp32 arithmetic - issue a third fewer, and the shared-memory tile keeps the bytes in flight - out of registers. + layers' inputs): 16384 tokens / 3 contexts 136 (the three-call chain it + replaced 528). Numerics: the conv result is rounded to bf16 before the norm reads it, as the unfused path stores it. Conv taps: fp32 FMAs from zero in tap order (`fma.rn.f32x2` on channel pairs). SiLU: `acc * rcp(1 + 2^(-acc * - log2e))` with `ex2.approx.ftz` and `rcp.approx.ftz` -- the Triton pass's - `ex2.approx` / `div.full` values wherever `1 + e <= 2^126`; below - acc = -87.3 it returns -0 where those return a value under 1e-36. q/k - norm: the 128-channel fp32 sum of squares as each thread's pairs plus a - 32-lane butterfly, `sqrt.approx.ftz`, `div.full`. v, g, beta and the pool - are bitwise the Triton pass's; q/k one bf16 ulp on ~3e-6 of elements - (the sum's order). Against the unfused path (entry test, every cell): + log2e))` with `ex2.approx.ftz` and `rcp.approx.ftz` -- a Triton + `acc / (1 + exp(-acc))`'s `ex2.approx` / `div.full` values wherever + `1 + e <= 2^126`; below acc = -87.3 it returns -0 where those return a + value under 1e-36. q/k norm: the 128-channel fp32 sum of squares as each + thread's pairs plus a 32-lane butterfly, `sqrt.approx.ftz`, `div.full`. + Against the unfused path (entry test, every cell): q/k/v within one bf16 ulp on 1e-6..1e-4 of elements, g/beta within the fp32 band, conv pool bitwise. @@ -1113,7 +929,9 @@ class GdnPrefillConvPost(OpWrapper): 4]` in `x`'s dtype (any strides); `conv_states` `[slots, qkv_dim, 3]` in `x`'s dtype (any strides); int32 `query_start_loc`/`cache_indices`, bool `has_initial_state`. Loud (`__call__`): `head_k_dim != head_v_dim`, conv - width other than 4. + width other than 4, and what the conv pass cannot take: an `x` not bf16, + heads other than 128 channels, a row of `x` not 16-byte aligned, a weight + or pool stride from 2^31 elements. Not checked: `cache_indices` within the pool and distinct; `query_start_loc` non-decreasing and ending at `T`. @@ -1121,8 +939,8 @@ class GdnPrefillConvPost(OpWrapper): Not certified: fp16/fp32 activations, other head geometries, bias. Hard limits - element offsets int32 unless the call needs int64 (chosen on the host - from T and the strides) not probed beyond 16384 tokens. + element offsets 64-bit for the rows of x and the outputs; weight and + pool strides int32 (refused from 2^31) not probed beyond 16384 tokens. """ def __call__( diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py index c204b6d2a508..4e1c6efaf56d 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py @@ -1206,19 +1206,6 @@ def _classify_generation_requests( else: generation_requests.append(request) extend_requests += extend_dummy_requests - # The engine no longer builds first-draft requests; one that still carries the - # flag is refused on an mRoPE model rather than packed without a delta slot. - first_draft_requests = [ - request - for request in scheduled_requests.generation_requests - if getattr(request, "py_is_first_draft", False) - ] - if host.use_mrope and first_draft_requests: - raise NotImplementedError( - "First-draft requests on an mRoPE model are not supported: they add no " - "mRoPE delta read slot, so every later generation request would read " - "another request's delta." - ) def _prepare_helix_state( self, host: HostInputs, kv_cache_manager: KVCacheManager | KVCacheManagerV2 | None diff --git a/tests/unittest/_torch/executor/test_mrope_first_draft_refusal.py b/tests/unittest/_torch/executor/test_mrope_first_draft_refusal.py deleted file mode 100644 index 26947a2e5782..000000000000 --- a/tests/unittest/_torch/executor/test_mrope_first_draft_refusal.py +++ /dev/null @@ -1,54 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""A first-draft request in a speculative batch of an mRoPE model is refused. - -The first-draft packing adds neither an mRoPE position span nor a delta read slot, and the -attention kernel indexes the deltas by generation batch index, so every generation request -packed after it would read another request's delta. The engine builds such a request nowhere -today; the runner says so loudly instead of misaligning the slots if one ever appears. -""" - -from unittest.mock import Mock - -import pytest -import test_pytorch_model_engine as engine_tests # module import: its TestCase is not collected - -from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests -from tensorrt_llm.llmapi import MTPDecodingConfig - - -def _prepare(first_draft: bool) -> None: - case = engine_tests.PyTorchModelEngineTestCase() - model_engine, kv_cache_manager, attn_metadata = case._setup_mrope_engine( - spec_config=MTPDecodingConfig(max_draft_len=1) - ) - try: - # A generation request with an mRoPE delta, then a plain one after it. - first = case._make_mrope_gen_request(4, 1, 0, 10) - first.py_is_first_draft = first_draft - plain = case._make_mrope_gen_request(6, 2, 1, 20) - scheduled_requests = ScheduledRequests() - scheduled_requests.context_requests_last_chunk = [] - scheduled_requests.generation_requests = [first, plain] - model_engine._runner._inputs.prepare_inputs( - scheduled_requests=scheduled_requests, - kv_cache_manager=kv_cache_manager, - attn_metadata=attn_metadata, - spec_metadata=Mock( - _force_non_greedy_for_capture=False, context_prompt_lookahead_tokens=None - ), - enable_spec_decode=True, - runtime_draft_len=1, - is_dummy=False, - ) - finally: - kv_cache_manager.shutdown() - - -def test_first_draft_request_on_an_mrope_model_is_refused() -> None: - with pytest.raises(NotImplementedError, match="First-draft requests on an mRoPE model"): - _prepare(first_draft=True) - - -def test_the_same_batch_without_a_first_draft_request_is_prepared() -> None: - _prepare(first_draft=False) diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 6e546e3a9a38..f7cc60f0a500 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -3293,37 +3293,6 @@ def test_prepare_inputs_mrope_verify_as_context_alone_emits_no_read_slots( [0]) kv_cache_manager.shutdown() - def _make_first_draft_batch(self): - """A first-draft request ahead of two plain generation requests that - carry MRoPE deltas.""" - first_draft = self._make_mrope_gen_request(4, 1, 0, 10) - first_draft.py_is_first_draft = True - first_draft.py_num_accepted_draft_tokens = 0 - scheduled_requests = ScheduledRequests() - scheduled_requests.context_requests_last_chunk = [] - scheduled_requests.generation_requests = [ - first_draft, - self._make_mrope_gen_request(6, 2, 1, 20), - self._make_mrope_gen_request(5, 3, 2, 30), - ] - return scheduled_requests - - def test_prepare_inputs_first_draft_requests_without_mrope_are_served( - self) -> None: - """Without MRoPE a first-draft request packs as a plain generation row.""" - model_engine, kv_cache_manager, attn_metadata = self._setup_mrope_engine( - spec_config=MTPDecodingConfig(max_draft_len=1)) - _set_runner_config(model_engine._runner, use_mrope=False) - result = self._prepare_mrope_spec_inputs(model_engine, kv_cache_manager, - attn_metadata, - self._make_first_draft_batch()) - - # The first-draft row is [last two prompt tokens + draft slot] wide. - self.assertEqual(result["position_ids"].view(-1).cpu().tolist(), - [2, 3, 5, 4]) - self.assertNotIn("mrope_delta_read_seq_slots", result) - kv_cache_manager.shutdown() - @skip_pre_blackwell_unittest def test_ced_warmup_passes_dummy_context_requests(self) -> None: for warmup in (False, True): diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py index 3d6a89a2c20b..5bfb34d165cc 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py @@ -64,22 +64,14 @@ def __call__(self, selected_tokens, *args, **kwargs): "_row_bits", "_rowtile_plan_kernel", "_rowtile_lone_kernel", - "_rowtile_attention_kernel", + "_rowtile_attention_cute", ) -#: The row-tile pass on SM103 tensor cores, which stands in for `_rowtile_attention_kernel` -#: on eligible launches (bf16 q and K/V, D = 256). -_CUTE_PASS = "_rowtile_attention_cute" - - -def _expected_kernels(args) -> list: - cute = _module._cute_eligible(args["q"], args["k_cache"], args["v_cache"], 128, 64) - return [*_KERNELS[:3], _CUTE_PASS if cute else _KERNELS[3]] @pytest.fixture def launches(monkeypatch): log: list = [] - for name in (*_KERNELS, _CUTE_PASS): + for name in _KERNELS: inner = getattr(_module, name) monkeypatch.setattr(_module, name, _LaunchRecorder(inner, name, log)) return log @@ -217,7 +209,7 @@ def test_certified_cells(cell, certifying, launches) -> None: # noqa: F811 assert bool(_is_prefix_packed(selected).all()) assert torch.equal((selected >= 0).sum(1), _bound(args["query_positions"], spec)) out = op(**args) - assert [name for name, _ in launches] == _expected_kernels(args) + assert [name for name, _ in launches] == list(_KERNELS) groups = math.ceil(rows / 10) assert launches[0][1] == (rows,) and launches[1][1] == (groups,) assert launches[2][1] == (min(rows, 256), spec["kv_heads"]) @@ -323,7 +315,7 @@ def test_zero_rows_return_an_empty_output() -> None: lambda a: {**a, "query_positions": None, "compress_ratio": None}, "the row-tile kernel needs the prefill pair", ), - "float32_query": (lambda a: {**a, "q": a["q"].float()}, "QSA queries must use BF16 or FP16"), + "float32_query": (lambda a: {**a, "q": a["q"].float()}, "QSA queries must use BF16"), "int64_selection": ( lambda a: {**a, "selected_tokens": a["selected_tokens"].long()}, "QSA selected tokens must be a two-dimensional int32 tensor", @@ -402,8 +394,8 @@ def test_packed_walk_is_the_resident_union(segments, shared_scores, monkeypatch) if not shared_scores: args["block_table"][:, 3] = -1 # an unmapped page inside every request's span log: list = [] - for name in ("_rowtile_attention_kernel", _CUTE_PASS): - monkeypatch.setattr(_module, name, _PlanRecorder(getattr(_module, name), log)) + name = "_rowtile_attention_cute" + monkeypatch.setattr(_module, name, _PlanRecorder(getattr(_module, name), log)) out = op(**args) _compare(out, op.reference(**args)) plan = log[0].cpu() @@ -415,54 +407,3 @@ def test_packed_walk_is_the_resident_union(segments, shared_scores, monkeypatch) assert torch.equal(plan[taken, 0], (plan[taken, 1] + 63) // 64) if shared_scores: assert int(plan[taken, 0].sum()) * 2 < int(dense[taken].sum()) - - -def _triton_bits(args): - """The Triton bitmap pass (`_rowtile_bits_kernel`) into a zeroed bitmap.""" - selected, positions = args["selected_tokens"], args["query_positions"].contiguous() - tpb, cr, width = args["tokens_per_block"], args["compress_ratio"], args["block_table"].shape[1] - rows, final_topk = selected.shape - num_words = (width * tpb + 63) // 64 - bits = torch.zeros((rows, num_words), dtype=torch.int64, device=DEV) - info = torch.empty((rows, 2), dtype=torch.int32, device=DEV) - _module._rowtile_bits_kernel[(rows,)]( - selected, positions, bits, info, selected.stride(0), selected.stride(1), bits.stride(0), - TOPK=final_topk, BLOCK_TOPK=(final_topk - cr + 1) // cr, COMPRESS_RATIO=cr, - TOKENS_PER_BLOCK=tpb, PAGE_TABLE_WIDTH=width, BOUND_TILE=64, BLOCK_C=1024, num_warps=4, - ) # fmt: skip - return bits, info - - -@pytest.mark.parametrize( - "order, holes, positions_dtype", - [ - ("emitted", False, torch.int32), - ("shuffled", True, torch.int64), - ("emitted", True, torch.int64), - ], - ids=["emitted", "shuffled-with-holes", "holes-int64-positions"], -) -def test_bitmap_pass_is_the_triton_pass_bitwise(order, holes, positions_dtype) -> None: - """The CuTe bitmap pass writes the Triton pass's `bits` and `row_info` word for word: - on the emitted token order, on columns shuffled within each row (any order is - correct; order only changes the shared-memory flushes), and with an all-padding row - and tokens past the table.""" - args = _build(_with([(0, 900), (907, 2200), (5000, 37)], slots=16), seed=8100) - args["query_positions"] = args["query_positions"].to(positions_dtype) - selected = args["selected_tokens"] - if order == "shuffled": - perm = torch.argsort(torch.rand(selected.shape, device=DEV), dim=1) - selected = torch.gather(selected, 1, perm) - if holes: - selected = selected.clone() - selected[1] = -1 - selected[5, :8] = args["block_table"].shape[1] * args["tokens_per_block"] + 3 - args["selected_tokens"] = selected.contiguous() - want_bits, want_info = _triton_bits(args) - width = args["block_table"].shape[1] - got_bits, got_info = _module._row_bits( - args["selected_tokens"], args["query_positions"], args["compress_ratio"], - args["tokens_per_block"], width, want_bits.shape[1], - ) # fmt: skip - assert torch.equal(got_info, want_info) - assert torch.equal(got_bits, want_bits) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_first.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_first.py deleted file mode 100644 index 9126b995746e..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_first.py +++ /dev/null @@ -1,450 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the qwen4_exp_hc_decode_first entry. - -Cells against the float64 reference: `hc_out` within one ulp (as -`qwen4_exp_hc_inject_norm`), `part` within `2^-16` of the reference (relative -to its largest element, plus the same relative term), `ss` to fp32 accuracy. -The MoE-combine cells widen `hc_out`'s band by the combine's own rounding and -check `part` / `ss` against the projection of the `hc_out` the op stored (see -`_compare_moe`). Beyond the cells: wrong formulas land outside, bitwise -agreement with the combine + this op composition, the weight-orientation -guard, determinism 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_hc_decode_first import ( # noqa: E501 - COMBINE_CONFIG, - FOLD_CONFIG, - launch, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_decode_first import ( # noqa: E501 - qwen4_exp_hc_decode_first as op, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp4_block_scale_moe_renorm_permuted import ( # noqa: E501 - finalize, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fused_sigmoid_gate_mul_add import ( # noqa: E501 - fused_sigmoid_gate_mul_add, -) - -__extra_import_path__ = ["../../../.."] -from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name - -assert torch.cuda.is_available(), "qwen4_exp_hc_decode_first requires a CUDA device" - -DEV = "cuda" -HC_COUNT, HIDDEN, RANK = 4, 2560, 320 -# This checkpoint's layer-0 hyper-connection element statistics. -DOWN_RMS, INJECT_RMS, GAMMA_RMS, GAMMA_ABSMAX = 0.0296, 0.0107, 0.477, 6.625 -PART_EPS = 2.0**-16 - - -def _band(got, want, eps) -> float: - """Worst element's distance as a fraction of `rtol = eps`, `atol = eps * max|want|`.""" - w = want.double() - return ((got.double() - w).abs() / (eps * w.abs().max().item() + eps * w.abs())).max().item() - - -def _ulp(got, want) -> float: - eps, tiny = torch.finfo(want.dtype).eps, torch.finfo(want.dtype).tiny - w = want.double() - return ((got.double() - w).abs() / (eps * w.abs() + tiny)).max().item() - - -def _compare(out, ref) -> None: - """`hc_out` one ulp: an fp32 accumulate rounded once. `part`: a 2560-term - fp32 dot per stream, so ~2^-20 relative to its scale is the expected error; - the gate is 2^-16. `ss`: a 2560-term fp32 sum of squares.""" - (hc, part, ss), (ref_hc, ref_part, ref_ss) = out, ref - assert hc.dtype == ref_hc.dtype and hc.shape == ref_hc.shape - assert part.dtype == torch.float32 and part.shape == ref_part.shape and part.stride(2) == 1 - assert ss.dtype == torch.float32 and ss.shape == ref_ss.shape - if not hc.numel(): - return - assert _ulp(hc, ref_hc) <= 1.0, f"hc_out {_ulp(hc, ref_hc):.4g} ulp" - assert _band(part, ref_part, PART_EPS) <= 1.0, f"part {_band(part, ref_part, PART_EPS):.4g}x" - assert torch.allclose(ss, ref_ss, rtol=1e-5, atol=0), "ss" - - -def _combine(block_out, shared, gate_logits) -> torch.Tensor: - """The MoE block output in float64, rounded once to `block_out`'s dtype.""" - s = torch.sigmoid(gate_logits.double().reshape(-1, 1)) - return (block_out.double() + s * shared.double()).to(block_out.dtype) - - -def _routed(block_out, topk_weights, topk_rows) -> torch.Tensor: - """The routed row: `block_out` itself, or with the fold the float64 - top-k weighted sum of the permuted rows rounded once to their dtype.""" - if topk_rows is None: - return block_out - valid = (topk_rows >= 0).unsqueeze(-1) - rows = block_out.double()[topk_rows.clamp(min=0).long()] - return ( - torch.where(valid, rows * topk_weights.double().unsqueeze(-1), 0.0) - .sum(1) - .to(block_out.dtype) - ) - - -def _hc_band_ratio( - hc, ref_hc, block_out, inject, shared, gate, topk_weights=None, topk_rows=None -) -> float: - """`hc_out` against the reference: one ulp of `hc_out` plus one ulp of the - combine times `|inject|` (the op's fp32 combine and the reference's float64 - one can round to neighbouring values at a midpoint); with the fold, one ulp - of the routed row times `|inject|` more (the op's fp32 FMA finalize and the - reference's exact sum can round apart the same way).""" - feps, tiny = torch.finfo(ref_hc.dtype).eps, torch.finfo(ref_hc.dtype).tiny - routed = _routed(block_out, topk_weights, topk_rows) - b = _combine(routed, shared, gate).double().abs() - if topk_rows is not None: - b = b + routed.double().abs() - band = feps * ref_hc.double().abs() + ( - feps * b.unsqueeze(-2) * inject.double().abs().unsqueeze(-1) - ).flatten(-2) - return ((hc.double() - ref_hc.double()).abs() / (band + tiny)).max().item() - - -def _compare_moe(args, out, ref) -> None: - """MoE-combine cells: `hc_out` in `_hc_band_ratio`'s band of the float64 - result; `part` and `ss` as `_compare` against `reference(hc_out, None, - None, ...)` -- the projection of the residual the op stored.""" - hyper_input, block_out, inject, gamma, first, c, n, shared, gate = args[:9] - topk_weights, topk_rows = args[9:] if len(args) > 9 else (None, None) - (hc, part, ss), (ref_hc, _, _) = out, ref - assert hc.dtype == ref_hc.dtype and hc.shape == ref_hc.shape - if not hc.numel(): - return - ratio = _hc_band_ratio(hc, ref_hc, block_out, inject, shared, gate, topk_weights, topk_rows) - assert ratio <= 1.0, f"hc_out {ratio:.4g}x the band" - _compare((hc, part, ss), op.reference(hc, None, None, gamma, first, c, n)) - - -def _build(spec, seed): - rows = spec["rows"] - hc_count = spec.get("hc_count", HC_COUNT) - hidden = spec.get("hidden", HIDDEN) - dtype = spec.get("dtype", torch.bfloat16) - width = hc_count * hidden - n_cols = spec.get("n_cols", RANK + hc_count) - g = torch.Generator(device=DEV).manual_seed(seed) - - def normal(*shape, scale): - return torch.randn(*shape, generator=g, device=DEV) * scale - - hyper_input = normal(rows, width, scale=3.0).to(dtype) - if spec.get("layout") == "row_slice": - buf = normal(2 * rows, width, scale=1.0).to(dtype) - buf[::2] = hyper_input - hyper_input = buf[::2] - gamma = normal(width, scale=GAMMA_RMS) - gamma[0], gamma[width // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX - gamma = gamma.to(dtype) - n_rows = spec.get("n_rows", -(-n_cols // 128) * 128) - weight = torch.zeros(n_rows, width, device=DEV) - weight[:n_cols] = normal(n_cols, width, scale=DOWN_RMS) - weight[n_cols - hc_count : n_cols] *= INJECT_RMS / DOWN_RMS - first = weight.to(dtype).t() - block_out = inject = shared = gate_logits = None - if spec.get("inject", True): - block_out = normal(rows, hidden, scale=2.0).to(dtype) - inject = (2.0 * torch.sigmoid(normal(rows, hc_count, scale=1.0))).to(dtype) - if spec.get("moe", False): - shared = normal(rows, hidden, scale=2.0).to(dtype) - gate_logits = normal(rows, 1, scale=2.0).to(dtype) - if spec.get("gate") == "flat": - gate_logits = gate_logits.reshape(rows) - top_k = spec.get("fold", 0) - if not top_k: - return hyper_input, block_out, inject, gamma, first, hc_count, n_cols, shared, gate_logits - # The runner's unfinalized output: padded permuted rows (garbage where no - # slot maps), a distinct row per slot, softmax-like bf16 weights. - n_perm = rows * top_k + 37 - block_out = normal(n_perm, hidden, scale=2.0).to(dtype) - order = torch.randperm(n_perm, generator=g, device=DEV)[: rows * top_k] - topk_rows = order.reshape(rows, top_k).to(torch.int32) - if spec.get("drop"): - topk_rows[::2, -3:] = -1 - topk_weights = torch.softmax(normal(rows, top_k, scale=1.0), -1).to(dtype) - return ( - hyper_input, - block_out, - inject, - gamma, - first, - hc_count, - n_cols, - shared, - gate_logits, - topk_weights, - topk_rows, - ) - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name - certifying(op) - args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) - before = [a.clone() if isinstance(a, torch.Tensor) else a for a in args] - out = op(*args) - for a, b in zip(args, before): - assert not isinstance(a, torch.Tensor) or torch.equal(a, b), "every input is read-only" - if args[1] is None: - assert out[0] is args[0], "norm-only returns hyper_input itself" - if args[7] is None: - _compare(out, op.reference(*args)) - else: - _compare_moe(args, out, op.reference(*args)) - - -def test_measured_distance_from_float64() -> None: - for rows in (1, 16): - args = _build(dict(rows=rows), seed=rows) - (hc, part, ss), (rh, rp, rs) = op(*args), op.reference(*args) - print( - f"\n rows {rows}: hc {_ulp(hc, rh):.3g} ulp, part {_band(part, rp, PART_EPS):.4f}x, " - f"ss rel {((ss - rs).abs() / rs).max().item():.3g}" - ) - - -def test_deterministic() -> None: - args = _build(dict(rows=4), seed=3) - for a, b in zip(op(*args), op(*args)): - assert torch.equal(a, b) - - -def test_wrong_variants_land_outside_the_band() -> None: - """Each mis-reading keeps the shapes; each lands > 5x outside. The streams - carry different scales so a stream mix-up is visible.""" - args = list(_build(dict(rows=8), seed=5)) - scales = torch.tensor([0.5, 1.0, 2.0, 4.0], device=DEV).repeat_interleave(HIDDEN) - args[0] = (args[0].float() * scales).to(torch.bfloat16) - hyper_input, block_out, inject, gamma, first, c, n, _, _ = args - want_hc, want_part, _ = op.reference(*args) - pre = hyper_input.double().unflatten(-1, (c, HIDDEN)) - hc = want_hc.double().unflatten(-1, (c, HIDDEN)) - g = (1.0 + gamma.double()).unflatten(-1, (c, HIDDEN)) - w = first[:, :n].double().unflatten(0, (c, HIDDEN)) - - def proj(y): - return torch.einsum("tsi,sin->stn", y, w) - - variants = { - "w instead of 1 + w": proj(hc * (g - 1.0)), - "pre-inject stream": proj(pre * g), - "no (1 + w)": proj(hc), - "streams summed into every partial": proj(hc * g).sum(0, keepdim=True).expand(c, -1, -1), - } - for name, value in variants.items(): - ratio = _band(value.float(), want_part, PART_EPS) - print(f"\n part {name:<36} {ratio:10.4g}x") - assert ratio > 5.0, name - stream_minor = pre.flatten(-2) + ( - block_out.double().unsqueeze(-1) * inject.double().unsqueeze(-2) - ).flatten(-2) - ratio = _ulp(stream_minor.to(torch.bfloat16), want_hc) - print(f" hc_out stream-minor inject {ratio:10.4g} ulp") - assert ratio > 5.0 - - -def test_guard_refuses_a_row_major_weight() -> None: - """A contiguous `[width, N]` weight is read with rows and columns swapped - and returns plausible numbers unguarded.""" - args = list(_build(dict(rows=2), seed=6)) - want = op.reference(*args)[1] - args[4] = args[4].contiguous() - assert _band(op(*args)[1], want, PART_EPS) > 5.0, "a row-major weight should answer wrongly" - with validating(op), pytest.raises(AssertionError): - op(*args) - - -def test_cuda_graph_capture_and_replay() -> None: - args = list(_build(dict(rows=1), seed=8)) - static_in = args[0] - 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): - graph_out = op(*args) - fresh = _build(dict(rows=1), seed=18)[0] - static_in.copy_(fresh) - graph.replay() - torch.cuda.synchronize() - for a, b in zip(graph_out, op(fresh, *args[1:])): - assert torch.equal(a, b) - - -@pytest.mark.parametrize("rows,n_cols", [(1, RANK + HC_COUNT), (1, RANK), (16, RANK + HC_COUNT)]) -def test_moe_combine_matches_the_composition_it_replaces(rows, n_cols) -> None: - """Against `fused_sigmoid_gate_mul_add` (in place on the routed output, as - the forward called it) and then this op without the combine, at the tile - the combine path launches with: bitwise. (The default tile sums `part` in - another fp32 order; the cells bound both against the reference.)""" - args = _build(dict(rows=rows, moe=True, n_cols=n_cols), seed=200 + rows) - hyper_input, block_out, inject, gamma, first, c, n, shared, gate = args - fused = op(*args) - combined = fused_sigmoid_gate_mul_add(block_out.clone(), gate, shared) - hc, part, ss = launch(hyper_input, combined, inject, gamma, first, c, n, config=COMBINE_CONFIG) - for name, a, b in zip(("hc_out", "part", "ss"), fused, (hc, part[:, :, :n], ss)): - assert torch.equal(a, b), name - default = op(hyper_input, combined, inject, gamma, first, c, n) - assert torch.equal(fused[0], default[0]), "hc_out does not depend on the tile" - print( - f"\n rows {rows} cols {n_cols}: vs the default tile, part bitwise " - f"{(fused[1] == default[1]).float().mean().item():.4f} " - f"({_band(fused[1], default[1], PART_EPS):.3g}x the band apart), ss bitwise " - f"{(fused[2] == default[2]).float().mean().item():.4f}" - ) - - -def test_moe_wrong_variants_land_outside_the_band() -> None: - """Each mis-reading of the combine lands > 5x outside `_hc_band_ratio`.""" - args = _build(dict(rows=8, moe=True), seed=11) - hyper_input, block_out, inject, gamma, first, c, n, shared, gate = args - want_hc = op.reference(*args)[0] - r, sh, s = block_out.double(), shared.double(), torch.sigmoid(gate.double()) - variants = { - "no sigmoid": r + gate.double() * sh, - "gate on the routed half": s * r + sh, - "gate dropped": r + sh, - "shared dropped": r, - } - x = hyper_input.double().unflatten(-1, (c, HIDDEN)) - for name, blk in variants.items(): - wrong = (x + blk.unsqueeze(-2) * inject.double().unsqueeze(-1)).flatten(-2) - ratio = _hc_band_ratio(wrong.to(torch.bfloat16), want_hc, block_out, inject, shared, gate) - print(f"\n hc_out {name:<26} {ratio:10.4g}x the band") - assert ratio > 5.0, name - - -def test_guard_refuses_a_partial_combine() -> None: - args = list(_build(dict(rows=2, moe=True), seed=12)) - hyper_input, block_out, inject, gamma, first, c, n, shared, gate = args - with validating(op), pytest.raises(AssertionError): - op(hyper_input, block_out, inject, gamma, first, c, n, shared, None) - with validating(op), pytest.raises(AssertionError): - op(hyper_input, None, None, gamma, first, c, n, shared, gate) - - -def test_moe_cuda_graph_capture_and_replay() -> None: - args = list(_build(dict(rows=1, moe=True), seed=13)) - 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): - graph_out = op(*args) - fresh = _build(dict(rows=1, moe=True), seed=23) - for i in (0, 1, 7, 8): - args[i].copy_(fresh[i]) - graph.replay() - torch.cuda.synchronize() - for a, b in zip(graph_out, op(*args)): - assert torch.equal(a, b) - - -@pytest.mark.parametrize( - "rows,n_cols,drop", - [ - (1, RANK + HC_COUNT, False), - (1, RANK, False), - (4, RANK + HC_COUNT, True), - (16, RANK + HC_COUNT, False), - ], -) -def test_fold_matches_finalize_then_combine(rows, n_cols, drop) -> None: - """The fold against the runner's finalize (emulated exactly: fp32 FMAs in - slot order, one rounding; the MoE entry's test pins that emulation to the - runner bitwise) followed by this op's combine path at the same tile: - `hc_out` and `ss` bitwise; `part` bitwise from 2 rows, and at 1 row (where - the two Triton specializations reduce the dot in a different fp32 order) - within 2^-22 of it relative to its scale. And `hc_out` bitwise the combine - path at its own tile.""" - args = _build(dict(rows=rows, moe=True, fold=10, drop=drop, n_cols=n_cols), seed=300 + rows) - hyper_input, perm, inject, gamma, first, c, n, shared, gate, topk_w, topk_rows = args - fused = op(*args) - routed = finalize(perm, topk_w, topk_rows) - hc, part, ss = launch( - hyper_input, routed, inject, gamma, first, c, n, shared, gate, config=FOLD_CONFIG - ) - part = part[:, :, :n] - assert torch.equal(fused[0], hc), "hc_out" - assert torch.equal(fused[2], ss), "ss" - if rows > 1: - assert torch.equal(fused[1], part), "part" - else: - assert _band(fused[1], part, 2.0**-22) <= 1.0, "part" - print( - f"\n rows {rows} cols {n_cols}: part vs finalize + combine bitwise " - f"{(fused[1] == part).float().mean().item():.4f} ({_band(fused[1], part, PART_EPS):.3g}x the band apart)" - ) - combined = op(hyper_input, routed, inject, gamma, first, c, n, shared, gate) - assert torch.equal(fused[0], combined[0]), "hc_out does not depend on the tile" - - -def test_fold_wrong_variants_land_outside_the_band() -> None: - """Each mis-reading of the finalize lands > 5x outside `_hc_band_ratio`.""" - args = _build(dict(rows=8, moe=True, fold=10), seed=17) - hyper_input, perm, inject, gamma, first, c, n, shared, gate, topk_w, topk_rows = args - want_hc = op.reference(*args)[0] - rows = perm.double()[topk_rows.long()] - w = topk_w.double().unsqueeze(-1) - s = torch.sigmoid(gate.double()).reshape(-1, 1) - variants = { - "unweighted sum": rows.sum(1), - "first slot only": rows[:, 0] * w[:, 0], - "weights by the wrong slot": (rows * w.flip(1)).sum(1), - "token-major rows (no index map)": (perm.double()[: 8 * 10].reshape(8, 10, -1) * w).sum(1), - } - x = hyper_input.double().unflatten(-1, (c, HIDDEN)) - for name, routed in variants.items(): - blk = routed.to(torch.bfloat16).double() + s * shared.double() - wrong = (x + blk.unsqueeze(-2) * inject.double().unsqueeze(-1)).flatten(-2) - ratio = _hc_band_ratio( - wrong.to(torch.bfloat16), want_hc, perm, inject, shared, gate, topk_w, topk_rows - ) - print(f"\n hc_out {name:<34} {ratio:10.4g}x the band") - assert ratio > 5.0, name - - -def test_guard_refuses_a_partial_fold() -> None: - args = list(_build(dict(rows=2, moe=True, fold=10), seed=19)) - hyper_input, perm, inject, gamma, first, c, n, shared, gate, topk_w, topk_rows = args - with validating(op), pytest.raises(AssertionError, match="needs the combine"): - op(hyper_input, perm, inject, gamma, first, c, n, None, None, topk_w, topk_rows) - with validating(op), pytest.raises(AssertionError, match="come together"): - op(hyper_input, perm, inject, gamma, first, c, n, shared, gate, None, topk_rows) - with validating(op), pytest.raises(AssertionError, match="int32"): - op(hyper_input, perm, inject, gamma, first, c, n, shared, gate, topk_w, topk_rows.long()) - - -def test_fold_cuda_graph_capture_and_replay() -> None: - args = list(_build(dict(rows=1, moe=True, fold=10), seed=29)) - 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): - graph_out = op(*args) - fresh = _build(dict(rows=1, moe=True, fold=10), seed=31) - for i in (0, 1, 2, 7, 8, 9, 10): - args[i].copy_(fresh[i]) - graph.replay() - torch.cuda.synchronize() - for a, b in zip(graph_out, op(*args)): - assert torch.equal(a, b) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py index 6a419915ba65..9a7ffbe1e12f 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py @@ -2,13 +2,12 @@ # SPDX-License-Identifier: Apache-2.0 """GPU test for the qwen4_exp_hc_decode_first_cute entry. -Cells against the Triton entry's float64 reference, in the Triton entry's bands: -`hc_out` within one ulp (plus the combine's and the finalize's own rounding at the -MoE sites), `part` within `2^-16` of the reference relative to its scale, `ss` to -fp32 accuracy. Beyond the cells: `hc_out` bitwise the Triton entry's (the same fp32 -arithmetic), `part` / `ss` inside the band of it; a stale-read check under CUDA-graph -replay with PDL, where a kernel that read its activations before -`griddepcontrol.wait` would see the previous replay's; determinism; the guards. +Cells against the entry's float64 reference: `hc_out` within one ulp (plus the +combine's and the finalize's own rounding at the MoE sites), `part` within `2^-16` of +the reference relative to its scale, `ss` to fp32 accuracy. Beyond the cells: a +stale-read check under CUDA-graph replay with PDL, where a kernel that read its +activations before `griddepcontrol.wait` would see the previous replay's; determinism; +the guards. """ import zlib @@ -16,9 +15,6 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_decode_first import ( # noqa: E501 - qwen4_exp_hc_decode_first as triton_op, -) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_decode_first_cute import ( # noqa: E501 qwen4_exp_hc_decode_first_cute as op, ) @@ -70,7 +66,7 @@ def _routed(block_out, topk_weights, topk_rows): def _hc_band_ratio( hc, ref_hc, block_out, inject, shared, gate, topk_weights=None, topk_rows=None ) -> float: - """The Triton entry's `hc_out` band at the MoE sites: one ulp of `hc_out`, plus one + """The `hc_out` band at the MoE sites: one ulp of `hc_out`, plus one ulp of the combine (and with the fold, of the routed row) times `|inject|`.""" feps, tiny = torch.finfo(ref_hc.dtype).eps, torch.finfo(ref_hc.dtype).tiny routed = _routed(block_out, topk_weights, topk_rows) @@ -94,7 +90,7 @@ def _compare_moe(args, out, ref) -> None: def _build(spec, seed): - """The Triton entry's test operands at one row.""" + """The test operands at one row.""" rows, hc_count, hidden, dtype = 1, HC_COUNT, HIDDEN, torch.bfloat16 width = hc_count * hidden n_cols = spec.get("n_cols", RANK + hc_count) @@ -166,32 +162,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by _compare_moe(args, out, op.reference(*args)) -@pytest.mark.parametrize( - "spec", - [ - dict(), - dict(moe=True), - dict(moe=True, fold=10), - dict(moe=True, fold=10, drop=True), - dict(inject=False), - ], - ids=["inject", "combine", "fold", "fold-drop", "norm-only"], -) -def test_matches_the_triton_entry(spec) -> None: - """`hc_out` bitwise (same fp32 arithmetic and rounding); `part` and `ss` reduce in - another order, inside the band of the Triton entry's values.""" - args = _build(spec, seed=4242) - hc, part, ss = op(*args) - t_hc, t_part, t_ss = triton_op(*args) - assert torch.equal(hc, t_hc), "hc_out" - assert _band(part, t_part, PART_EPS) <= 1.0, "part" - assert torch.allclose(ss, t_ss, rtol=1e-5, atol=0), "ss" - print( - f"\n {spec}: part bitwise {(part == t_part).float().mean().item():.3f} " - f"({_band(part, t_part, PART_EPS):.3g}x the band apart)" - ) - - def test_deterministic() -> None: args = _build(dict(moe=True, fold=10), seed=3) for a, b in zip(op(*args), op(*args)): diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix.py deleted file mode 100644 index 306e1049b8d7..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix.py +++ /dev/null @@ -1,157 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the qwen4_exp_hc_decode_mix entry: cells inside -`qwen4_exp_gated_residual`'s band (`rtol = eps`, `atol = eps * max|ref|` of the -activation dtype) against the float64 reference for both `mixed` and -`inject_out`, wrong formulas outside, the stride guards, graph capture.""" - -import zlib - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_decode_mix import ( # noqa: E501 - qwen4_exp_hc_decode_mix as op, -) - -__extra_import_path__ = ["../../../.."] -from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name - -assert torch.cuda.is_available(), "qwen4_exp_hc_decode_mix requires a CUDA device" - -DEV = "cuda" -HC_COUNT, HIDDEN, RANK, EPS = 4, 2560, 320, 1e-6 -UP_RMS, GAMMA_RMS, GAMMA_ABSMAX = 0.0681, 0.477, 6.625 - - -def _band(got, want) -> float: - eps = torch.finfo(want.dtype).eps - w = want.double() - return ((got.double() - w).abs() / (eps * w.abs().max().item() + eps * w.abs())).max().item() - - -def _compare(out, ref) -> None: - (mixed, inj), (ref_mixed, ref_inj) = out, ref - assert (inj is None) == (ref_inj is None) - for got, want in [(mixed, ref_mixed)] + ([] if inj is None else [(inj, ref_inj)]): - assert got.dtype == want.dtype and got.shape == want.shape and got.is_contiguous() - if got.numel(): - assert _band(got, want) <= 1.0, f"{_band(got, want):.4g}x the band" - - -def _build(spec, seed): - rows = spec["rows"] - hc_count = spec.get("hc_count", HC_COUNT) - hidden = spec.get("hidden", HIDDEN) - rank = spec.get("rank", RANK) - dtype = spec.get("dtype", torch.bfloat16) - with_inject = spec.get("with_inject", True) - width = hc_count * hidden - g = torch.Generator(device=DEV).manual_seed(seed) - - def normal(*shape, scale): - return torch.randn(*shape, generator=g, device=DEV) * scale - - hc_out = normal(rows, width, scale=3.0).to(dtype) - if spec.get("layout") == "row_slice": - buf = normal(2 * rows, width, scale=1.0).to(dtype) - buf[::2] = hc_out - hc_out = buf[::2] - ss = hc_out.double().unflatten(-1, (hc_count, hidden)).pow(2).sum(-1).t().float().contiguous() - gamma = normal(width, scale=GAMMA_RMS) - gamma[0], gamma[width // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX - gamma = gamma.to(dtype) - # Per-stream partials sized so the combined logits reach this checkpoint's ~|67|. - n_pad = -(-(rank + hc_count) // 4) * 4 - part = normal(hc_count, rows, n_pad + 4, scale=60.0)[:, :, :n_pad] - up = normal(width, rank, scale=UP_RMS).to(dtype).t() - return part, ss, up, hc_out, gamma, EPS, hc_count, with_inject - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name - certifying(op) - args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) - out = op(*args) - assert out[0].shape == (cell.spec["rows"], args[3].shape[1] // args[6]) - _compare(out, op.reference(*args)) - - -def test_measured_distance_from_float64() -> None: - for rows in (1, 16): - args = _build(dict(rows=rows), seed=rows) - (m, i), (rm, ri) = op(*args), op.reference(*args) - print(f"\n rows {rows}: mixed {_band(m, rm):.4f}x, inject {_band(i, ri):.4f}x the band") - - -def test_programs_agree_and_repeat() -> None: - args = _build(dict(rows=4), seed=3) - a, b = op(*args), op(*args) - assert torch.equal(a[0], b[0]) and torch.equal(a[1], b[1]) - - -def test_wrong_variants_land_outside_the_band() -> None: - part, ss, up, hc_out, gamma, eps, c, _ = args = _build(dict(rows=8), seed=7) - want_mixed, want_inj = op.reference(*args) - rstd = torch.rsqrt(ss.double().t() / HIDDEN + eps) - x = hc_out.double().unflatten(-1, (c, HIDDEN)) - g = (1.0 + gamma.double()).unflatten(-1, (c, HIDDEN)) - logits = torch.einsum("ts,stn->tn", rstd, part.double()) - raw_logits = part.double().sum(0) - - def mix(lg, normed, reduce="mean"): - gate = torch.sigmoid(torch.nn.functional.silu(lg[:, :RANK] / c) @ up.double()) - v = gate.unflatten(-1, (c, HIDDEN)) * normed - return v.mean(-2) if reduce == "mean" else v.sum(-2) - - normed = x * rstd.unsqueeze(-1) * g - variants = { - "sum instead of mean": mix(logits, normed, "sum"), - "gates the raw stream": mix(logits, x), - "w instead of 1 + w": mix(logits, x * rstd.unsqueeze(-1) * (g - 1.0)), - "no rstd on the logits": mix(raw_logits, normed), - "no /hc_count before silu": mix(logits * c, normed), - } - for name, value in variants.items(): - ratio = _band(value.to(torch.bfloat16), want_mixed) - print(f"\n mixed {name:<28} {ratio:10.4g}x the band") - assert ratio > 5.0, name - zi = logits[:, RANK : RANK + c] - for name, value in { - "inject missing the factor of 2": torch.sigmoid(zi / c), - "inject without rstd": 2.0 * torch.sigmoid(raw_logits[:, RANK : RANK + c] / c), - }.items(): - ratio = _band(value.to(torch.bfloat16), want_inj) - print(f" {name:<34} {ratio:10.4g}x the band") - assert ratio > 5.0, name - - -def test_guard_refuses_a_strided_ss() -> None: - """`ss` laid out `[T, C]` and viewed as `[C, T]` is read row-major and - answers plausibly wrong unguarded; the guard refuses it.""" - args = list(_build(dict(rows=8), seed=6)) - want = op.reference(*args)[0] - args[1] = args[1].t().contiguous().t() - assert _band(op(*args)[0], want) > 1.0, "a transposed ss should answer outside the band" - with validating(op), pytest.raises(AssertionError): - op(*args) - - -def test_cuda_graph_capture_and_replay() -> None: - args = list(_build(dict(rows=1), seed=8)) - 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): - graph_out = op(*args) - fresh = _build(dict(rows=1), seed=18) - for k in (0, 1, 3): - args[k].copy_(fresh[k]) - graph.replay() - torch.cuda.synchronize() - eager = op(*args) - assert torch.equal(graph_out[0], eager[0]) and torch.equal(graph_out[1], eager[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_hc_decode_mix_cute.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_cute.py index fc1a21f55ba8..93e21cd02dca 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_cute.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_cute.py @@ -1,18 +1,15 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """GPU test for the qwen4_exp_hc_decode_mix_cute entry (the plain build of the CuTe mix -kernel): cells against the Triton entry's float64 reference in its band (`mixed` and -`inject_out` in `rtol = eps`, `atol = eps * max|ref|`), agreement with the Triton entry -inside the same band, a stale-read check under CUDA-graph replay with PDL, the guards.""" +kernel): cells against the entry's float64 reference in its band (`mixed` and +`inject_out` in `rtol = eps`, `atol = eps * max|ref|`), a stale-read check under +CUDA-graph replay with PDL, 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.hc.qwen4_exp_hc_decode_mix import ( # noqa: E501 - qwen4_exp_hc_decode_mix as triton_op, -) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_decode_mix_fp4_cute import ( # noqa: E501 qwen4_exp_hc_decode_mix_cute as op, ) @@ -74,17 +71,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by _compare(out, op.reference(*args)) -@pytest.mark.parametrize("with_inject", [True, False]) -def test_matches_the_triton_entry(with_inject) -> None: - args = _build(dict(with_inject=with_inject), seed=77) - mixed, inj = op(*args) - t_mixed, t_inj = triton_op(*args) - assert _band(mixed, t_mixed) <= 1.0 - if with_inject: - assert _band(inj, t_inj) <= 1.0 - print(f"\n mixed bitwise the Triton entry's: {(mixed == t_mixed).float().mean().item():.4f}") - - def test_deterministic() -> None: args = _build(dict(), seed=3) for a, b in zip(op(*args), op(*args)): diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4.py deleted file mode 100644 index d760d66b80ab..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4.py +++ /dev/null @@ -1,190 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the qwen4_exp_hc_decode_mix_fp4 entry: `mixed` and `inject_out` -inside `qwen4_exp_hc_decode_mix`'s band against the float64 reference, the -NVFP4 outputs exact on the returned `mixed` (outside the reciprocal tie window) -and bitwise against `qwen4_exp_hc_decode_mix` + `torch.ops.trtllm.fp4_quantize`, -guards, graph capture.""" - -import zlib - -import pytest -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_decode_mix import ( # noqa: E501 - qwen4_exp_hc_decode_mix, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_decode_mix_fp4 import ( # noqa: E501 - nvfp4_near_tie, - nvfp4_reference, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_decode_mix_fp4 import ( # noqa: E501 - qwen4_exp_hc_decode_mix_fp4 as op, -) - -__extra_import_path__ = ["../../../.."] -from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name - -assert torch.cuda.is_available(), "qwen4_exp_hc_decode_mix_fp4 requires a CUDA device" - -DEV = "cuda" -HC_COUNT, HIDDEN, RANK, EPS = 4, 2560, 320, 1e-6 -UP_RMS, GAMMA_RMS, GAMMA_ABSMAX = 0.0681, 0.477, 6.625 - - -def _band(got, want) -> float: - eps = torch.finfo(want.dtype).eps - w = want.double() - return ((got.double() - w).abs() / (eps * w.abs().max().item() + eps * w.abs())).max().item() - - -def _unpack(xq: torch.Tensor) -> torch.Tensor: - return torch.stack([xq & 0xF, xq >> 4], -1).flatten(-2) - - -def _check_quant(mixed, xq, xsf, gs) -> None: - want_q, want_sf = nvfp4_reference(mixed, gs) - assert xq.dtype == torch.uint8 and xq.shape == want_q.shape and xq.is_contiguous() - assert xsf.dtype == torch.uint8 and xsf.shape == want_sf.shape - assert torch.equal(xsf, want_sf), f"{(xsf != want_sf).sum().item()} scale bytes differ" - stray = (_unpack(xq) != _unpack(want_q)) & ~nvfp4_near_tie(mixed, gs) - assert not stray.any(), f"{stray.sum().item()} e2m1 codes differ outside the tie window" - - -def _compare(out, ref, gs) -> None: - """`mixed` / `inject_out` in `rtol = eps`, `atol = eps * max|ref|` (the band - `qwen4_exp_hc_decode_mix` is certified in); NVFP4 exact on the returned - `mixed`, and equal to the reference's on 16-blocks whose `mixed` matches.""" - (mixed, inj, xq, xsf), (ref_mixed, ref_inj, ref_q, ref_sf) = out, ref - assert (inj is None) == (ref_inj is None) - for got, want in [(mixed, ref_mixed)] + ([] if inj is None else [(inj, ref_inj)]): - assert got.dtype == want.dtype and got.shape == want.shape and got.is_contiguous() - if got.numel(): - assert _band(got, want) <= 1.0, f"{_band(got, want):.4g}x the band" - if mixed.numel() == 0: - assert xq.numel() == 0 and xsf.numel() == 0 - return - _check_quant(mixed, xq, xsf, gs) - rows = mixed.shape[0] - same = (mixed.view(rows, -1, 16) == ref_mixed.view(rows, -1, 16)).all(-1) - assert torch.equal(xsf.view(rows, -1)[same], ref_sf.view(rows, -1)[same]) - q, rq = _unpack(xq).view(rows, -1, 16), _unpack(ref_q).view(rows, -1, 16) - tie = nvfp4_near_tie(ref_mixed, gs).view(rows, -1, 16) - assert not ((q != rq) & ~tie)[same].any() - - -def _build(spec, seed): - rows = spec["rows"] - hc_count = spec.get("hc_count", HC_COUNT) - hidden = spec.get("hidden", HIDDEN) - rank = spec.get("rank", RANK) - dtype = spec.get("dtype", torch.bfloat16) - with_inject = spec.get("with_inject", True) - width = hc_count * hidden - g = torch.Generator(device=DEV).manual_seed(seed) - - def normal(*shape, scale): - return torch.randn(*shape, generator=g, device=DEV) * scale - - hc_out = normal(rows, width, scale=3.0).to(dtype) - if spec.get("layout") == "row_slice": - buf = normal(2 * rows, width, scale=1.0).to(dtype) - buf[::2] = hc_out - hc_out = buf[::2] - ss = hc_out.double().unflatten(-1, (hc_count, hidden)).pow(2).sum(-1).t().float().contiguous() - gamma = normal(width, scale=GAMMA_RMS) - gamma[0], gamma[width // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX - gamma = gamma.to(dtype) - n_pad = -(-(rank + hc_count) // 4) * 4 - part = normal(hc_count, rows, n_pad + 4, scale=60.0)[:, :, :n_pad] - up = normal(width, rank, scale=UP_RMS).to(dtype).t() - # A calibrated activation global scale: 448 * 6 / amax of the mix (~|8|). - gs = torch.tensor([448.0 * 6.0 / 8.0], device=DEV, dtype=torch.float32) - return part, ss, up, hc_out, gamma, EPS, hc_count, with_inject, gs - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name - certifying(op) - args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) - out = op(*args) - assert out[0].shape == (cell.spec["rows"], args[3].shape[1] // args[6]) - _compare(out, op.reference(*args), args[8]) - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_bitwise_the_composition_it_replaces(cell) -> None: - """`qwen4_exp_hc_decode_mix` then `torch.ops.trtllm.fp4_quantize` (linear - scales). The NVFP4 outputs are bitwise `fp4_quantize` of the returned - `mixed`, ties included; `inject_out` is bitwise `qwen4_exp_hc_decode_mix`'s. - `mixed` is held to `qwen4_exp_hc_decode_mix`'s within the band: the - 16-column program tile reduces the up projection in another order than the - 8-column one, so a few elements in ten thousand round the other way - (measured: 0..5 of 10240 elements at 1..4 rows).""" - args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) - mixed, inj, xq, xsf = op(*args) - old_mixed, old_inj = qwen4_exp_hc_decode_mix(*args[:8]) - assert (inj is None and old_inj is None) or torch.equal(inj, old_inj) - if mixed.numel(): - diff = (mixed != old_mixed).sum().item() - print( - f"\n {mixed.dtype}: {diff} elements differ from decode_mix, " - f"{_band(mixed, old_mixed):.4f}x the band" - ) - assert _band(mixed, old_mixed) <= 1.0 - if mixed.shape[0] == 0: - return - old_q, old_sf = torch.ops.trtllm.fp4_quantize(mixed, args[8], 16, False, False) - assert torch.equal(xq, old_q.view(torch.uint8)) - assert torch.equal(xsf, old_sf.view(torch.uint8).reshape(-1)) - - -def test_wrong_recipes_fail_the_quant_check() -> None: - """Each plausible wrong layout or recipe is refused by `_check_quant`.""" - args = _build(dict(rows=16), seed=3) - mixed, _, xq, xsf = op(*args) - gs = args[8] - _check_quant(mixed, xq, xsf, gs) - swapped = ((xq & 0xF) << 4) | (xq >> 4) - shifted = torch.roll(xsf.view(16, -1), 1, dims=1).reshape(-1) - half_gs = nvfp4_reference(mixed, gs * 0.5) - for name, bad in ( - ("nibbles swapped", (mixed, swapped, xsf)), - ("scales one block off", (mixed, xq, shifted)), - ("global scale halved", (mixed,) + half_gs), - ): - with pytest.raises(AssertionError): - _check_quant(*bad, gs) - print(f"\n {name:<24} refused") - - -def test_guards() -> None: - args = list(_build(dict(rows=4), seed=4)) - with validating(op), pytest.raises(AssertionError): - op(*args[:8], torch.cat([args[8], args[8]])) - bad = list(_build(dict(rows=2, hidden=136, rank=32), seed=5)) - with validating(op), pytest.raises(AssertionError): - op(*bad) - - -def test_cuda_graph_capture_and_replay() -> None: - args = list(_build(dict(rows=1), seed=8)) - 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): - graph_out = op(*args) - fresh = _build(dict(rows=1), seed=18) - for k in (0, 1, 3): - args[k].copy_(fresh[k]) - args[8].mul_(0.5) - graph.replay() - torch.cuda.synchronize() - eager = op(*args) - for got, want in zip(graph_out, eager): - assert torch.equal(got, want) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py index 29724ea6a96a..761212e90c47 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py @@ -1,22 +1,19 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the qwen4_exp_hc_decode_mix_fp4_cute entry: cells against the Triton -entry's float64 reference in its band (`mixed` and `inject_out` in `rtol = eps`, +"""GPU test for the qwen4_exp_hc_decode_mix_fp4_cute entry: cells against the entry's +float64 reference in its band (`mixed` and `inject_out` in `rtol = eps`, `atol = eps * max|ref|`; the NVFP4 exact on the returned `mixed` outside the -reciprocal tie window), agreement with the Triton entry inside the same band, a -stale-read check under CUDA-graph replay with PDL, the guards.""" +reciprocal tie window), a stale-read check under CUDA-graph replay with PDL, 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.hc.qwen4_exp_hc_decode_mix_fp4 import ( # noqa: E501 - nvfp4_near_tie, - nvfp4_reference, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_decode_mix_fp4 import ( # noqa: E501 - qwen4_exp_hc_decode_mix_fp4 as triton_op, +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( + Fp4Quantize, + fp4_quantize, ) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_decode_mix_fp4_cute import ( # noqa: E501 qwen4_exp_hc_decode_mix_fp4_cute as op, @@ -32,6 +29,18 @@ UP_RMS, GAMMA_RMS, GAMMA_ABSMAX = 0.0681, 0.477, 6.625 +def nvfp4_reference(x, gs): + """`fp4_quantize`'s native-torch NVFP4 of `x` `[T, K]` with linear scales: + `(packed [T, K/2] uint8, sf [T*K/16] uint8)`.""" + return fp4_quantize.reference(x, gs, 16, False, False) + + +def nvfp4_near_tie(x, gs): + """`[T, K]` bool: elements within `fp4_quantize`'s tie window of an e2m1 midpoint, + where the kernels' approximate reciprocals may round either way.""" + return Fp4Quantize.near_tie(fp4_quantize.scaled(x, gs, 16)) + + def _band(got, want) -> float: eps = torch.finfo(want.dtype).eps w = want.double() @@ -97,19 +106,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by _check_quant(out[0], out[2], out[3], args[8]) -@pytest.mark.parametrize("with_inject", [True, False]) -def test_matches_the_triton_entry(with_inject) -> None: - """`mixed` / `inject_out` inside the band of the Triton entry's; the NVFP4 of each - is exact on its own `mixed`, so where the two `mixed` agree the codes agree.""" - args = _build(dict(with_inject=with_inject), seed=77) - mixed, inj, _, _ = op(*args) - t_mixed, t_inj, _, _ = triton_op(*args) - assert _band(mixed, t_mixed) <= 1.0 - if with_inject: - assert _band(inj, t_inj) <= 1.0 - print(f"\n mixed bitwise the Triton entry's: {(mixed == t_mixed).float().mean().item():.4f}") - - def test_deterministic() -> None: args = _build(dict(), seed=3) for a, b in zip(op(*args), op(*args)): diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_gate_mean.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_gate_mean.py deleted file mode 100644 index a6f4b23d4df7..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_gate_mean.py +++ /dev/null @@ -1,145 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the qwen4_exp_hc_gate_mean entry: cells in the band -`qwen4_exp_gated_residual`'s test derives for `mixed`, wrong formulas outside, -graph capture.""" - -import zlib - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501 - gate_ratio, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_gate_mean import ( # noqa: E501 - qwen4_exp_hc_gate_mean as op, -) - -__extra_import_path__ = ["../../../.."] -from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name - -assert torch.cuda.is_available(), "qwen4_exp_hc_gate_mean requires a CUDA device" - -DEV = "cuda" - - -def _compare(out, ref) -> None: - """`rtol = eps`, `atol = eps * max|ref|` (`gate_ratio <= 1`): `mixed` is a - mean of signed products, so near-zero elements are cancellations whose - error is set by the stream's scale -- the band `qwen4_exp_gated_residual` - derives for the same output.""" - assert out.dtype == ref.dtype and out.shape == ref.shape and out.is_contiguous() - if out.numel(): - assert gate_ratio(out, ref) <= 1.0, f"{gate_ratio(out, ref):.4g}x the band" - - -def _build(spec, seed): - hc_count, hidden = spec.get("hc_count", 4), spec.get("hidden", 2560) - dtype = spec.get("dtype", torch.bfloat16) - g = torch.Generator(device=DEV).manual_seed(seed) - rows, width = spec["rows"], hc_count * hidden - # Gate logits of this checkpoint's up projection span about +-8; the normed stream +-30. - gate_logits = (torch.randn(rows, width, generator=g, device=DEV) * 3).to(dtype) - normed = (torch.randn(rows, width, generator=g, device=DEV) * 6).to(dtype) - return gate_logits, normed, 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 — fixture by name - certifying(op) - args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) - mixed = op(*args) - assert mixed.shape == (cell.spec["rows"], cell.spec.get("hidden", 2560)) - _compare(mixed, op.reference(*args)) - - -def test_wrong_variants_land_outside_the_band() -> None: - gate_logits, normed, c = _build(dict(rows=128), seed=3) - want = op.reference(gate_logits, normed, c) - h = normed.shape[1] // c - g = torch.sigmoid(gate_logits.double()).unflatten(-1, (c, h)) - x = normed.double().unflatten(-1, (c, h)) - for name, value in ( - ("sum instead of mean", (g * x).sum(-2)), - ("no sigmoid", (gate_logits.double().unflatten(-1, (c, h)) * x).mean(-2)), - ("stream-minor packing", (g.transpose(-1, -2).reshape(-1, c, h) * x).mean(-2)), - ("ungated mean", x.mean(-2)), - ): - ratio = gate_ratio(value.to(want.dtype), want) - print(f"\n {name:<24} {ratio:10.4g}x the band") - assert ratio > 5.0, name - - -def test_guard_refuses_mismatched_shapes() -> None: - gate_logits, normed, c = _build(dict(rows=4), seed=4) - with validating(op), pytest.raises(AssertionError): - op(gate_logits[:2], normed, c) - - -def test_cuda_graph_capture_and_replay() -> None: - gate_logits, normed, c = _build(dict(rows=8), seed=5) - side = torch.cuda.Stream() - side.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(side): - op(gate_logits, normed, c) - torch.cuda.current_stream().wait_stream(side) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - g_mixed = op(gate_logits, normed, c) - fresh = _build(dict(rows=8), seed=6) - gate_logits.copy_(fresh[0]) - normed.copy_(fresh[1]) - graph.replay() - torch.cuda.synchronize() - assert torch.equal(g_mixed, op(fresh[0], fresh[1], c)) - - -@pytest.mark.parametrize("rows", [1, 33, 128, 2048, 8192]) -@pytest.mark.parametrize("with_inject", [True, False]) -def test_mixer_matches_the_gated_residual_reference(rows, with_inject) -> None: - """The mixer as the forward runs it -- `cublas_mm` over the zero-padded - `[down; inject]` (fp32 out), `qwen4_exp_hc_lowrank_act`, `cublas_mm` for - the up projection, this op -- against `qwen4_exp_gated_residual.reference` - (float64 from `normed`), in the band that entry certifies its own op in.""" - from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm - from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501 - qwen4_exp_gated_residual, - ) - from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_lowrank_act import ( # noqa: E501 - qwen4_exp_hc_lowrank_act, - ) - - hc, hidden, rank, eps = 4, 2560, 320, 1e-6 - width = hc * hidden - g = torch.Generator(device=DEV).manual_seed(rows) - - def normal(*shape, scale): - return (torch.randn(*shape, generator=g, device=DEV) * scale).to(torch.bfloat16) - - # This checkpoint's layer-0 hyper-connection element statistics. - down, up = normal(rank, width, scale=0.0296), normal(width, rank, scale=0.0681) - inject_w, gamma = normal(hc, width, scale=0.0107), normal(width, scale=0.477) - gamma[0], gamma[width // 2] = 6.625, -6.625 - x = torch.randn(rows, hc, hidden, generator=g, device=DEV) - x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) - normed = (x.flatten(-2) * (1.0 + gamma.float())).to(torch.bfloat16) - first = torch.cat([down, inject_w], 0) if with_inject else down - first = torch.cat( - [first, torch.zeros(384 - first.shape[0], width, device=DEV, dtype=first.dtype)] - ) - low, inject = qwen4_exp_hc_lowrank_act( - cublas_mm(normed, first.t(), None, torch.float32), rank, hc, with_inject, torch.bfloat16 - ) - mixed = op(cublas_mm(low, up.t()), normed, hc) - want = qwen4_exp_gated_residual.reference( - normed, down, up, inject_w if with_inject else None, hc - ) - _compare(mixed, want[0]) - if with_inject: - _compare(inject, want[1]) - old = qwen4_exp_gated_residual(normed, down, up, inject_w if with_inject else None, hc) - print( - f"\n rows {rows}: new mixed {gate_ratio(mixed, want[0]):.3f}x band, " - f"old op {gate_ratio(old[0], want[0]):.3f}x" - ) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py index 9fa943fee089..3521a59f384a 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py @@ -110,9 +110,6 @@ def test_mixer_matches_the_gated_residual_reference(rows, with_inject) -> None: 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, ) - from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_gate_mean import ( # noqa: E501 - qwen4_exp_hc_gate_mean, - ) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_lowrank_act import ( # noqa: E501 qwen4_exp_hc_lowrank_act, ) @@ -143,8 +140,4 @@ def normal(*shape, scale): normed, down, up, inject_w if with_inject else None, hc ) _compare(mixed, want[0]) - prev = qwen4_exp_hc_gate_mean(cublas_mm(low, up.t()), normed, hc) - print( - f"\n rows {rows}: fused {gate_ratio(mixed, want[0]):.3f}x band, " - f"cublas + gate_mean {gate_ratio(prev, want[0]):.3f}x" - ) + print(f"\n rows {rows}: fused {gate_ratio(mixed, want[0]):.3f}x band") diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4.py index a315944a3393..aa4b31487bdf 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4.py @@ -12,13 +12,13 @@ import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( + Fp4Quantize, + fp4_quantize, +) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501 gate_ratio, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_decode_mix_fp4 import ( # noqa: E501 - nvfp4_near_tie, - nvfp4_reference, -) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_up_gate_mean import ( # noqa: E501 qwen4_exp_hc_up_gate_mean, ) @@ -34,6 +34,18 @@ DEV = "cuda" +def nvfp4_reference(x, gs): + """`fp4_quantize`'s native-torch NVFP4 of `x` `[T, K]` with linear scales: + `(packed [T, K/2] uint8, sf [T*K/16] uint8)`.""" + return fp4_quantize.reference(x, gs, 16, False, False) + + +def nvfp4_near_tie(x, gs): + """`[T, K]` bool: elements within `fp4_quantize`'s tie window of an e2m1 midpoint, + where the kernels' approximate reciprocals may round either way.""" + return Fp4Quantize.near_tie(fp4_quantize.scaled(x, gs, 16)) + + def _unpack(xq: torch.Tensor) -> torch.Tensor: return torch.stack([xq & 0xF, xq >> 4], -1).flatten(-2) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_weight_path.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_weight_path.py index f692c332ac5c..cd6ffd05d6c1 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_weight_path.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_weight_path.py @@ -35,6 +35,9 @@ fp4_block_scale_moe_runner, quantize_nvfp4, ) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( + fp4_quantize, +) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 modeling as M, ) @@ -53,6 +56,10 @@ #: expert's. ISC2 = 64.0 / (448 * 6) E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) +#: The routed runners read activation scales as a linear buffer, not the 128x4-swizzled +#: one of a dense NVFP4 GEMM; byte counts match at multiples of 128 rows, so the wrong +#: one passes silently. +SF_LINEAR = False #: Tokens of the two routings: one per expert, and a weighted top-10. ONE_HOT_TOKENS = E @@ -170,7 +177,7 @@ def _cutedsl_routed(layer, x, ids, weights) -> torch.Tensor: """The CuteDSL branch of `_moe_block`, from the quantization to the unpermute.""" sc = layer.scalars tokens = x.shape[0] - xq, xsf = M.fp4_quantize(x, sc["global_scale"], M._SF_VEC, False, M._SF_LINEAR) + xq, xsf = fp4_quantize(x, sc["global_scale"], M._SF_VEC, False, SF_LINEAR) tile_expert, tile_limit, expanded_to_permuted, permuted_to_expanded, _, live_tiles = M.moe_sort( token_selected_experts=ids, token_final_scales=weights, @@ -226,7 +233,7 @@ def _cutedsl_routed(layer, x, ids, weights) -> torch.Tensor: def _trtllm_gen_routed(layer, x, ids, weights) -> torch.Tensor: """The trtllm-gen branch of `_moe_block`.""" sc = layer.scalars - xq, xsf = M.fp4_quantize(x, sc["global_scale"], M._SF_VEC, False, M._SF_LINEAR) + xq, xsf = fp4_quantize(x, sc["global_scale"], M._SF_VEC, False, SF_LINEAR) return fp4_block_scale_moe_runner( routing_logits=None, routing_bias=None, @@ -302,7 +309,7 @@ def _reference(backend: str, case: str, spread: bool) -> torch.Tensor: ck, layer = _checkpoint(spread), _layer(backend, spread) x, ids, weights = _problem(case) sc = layer.scalars - xq, xsf = M.fp4_quantize(x, sc["global_scale"], M._SF_VEC, False, M._SF_LINEAR) + xq, xsf = fp4_quantize(x, sc["global_scale"], M._SF_VEC, False, SF_LINEAR) x_dq = _dequant(xq.view(torch.uint8), xsf.view(x.shape[0], H // VEC)) * ISC1 if backend == "CUTEDSL": g2 = [float(sc["global_sf"])] * E diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_gate_conv_inject_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_gate_conv_inject_norm.py index 07a42aea5272..0cfbdf6de72c 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_gate_conv_inject_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_gate_conv_inject_norm.py @@ -82,9 +82,8 @@ def _stats(out: torch.Tensor, ref: torch.Tensor) -> tuple[float, float]: #: The per-element row-max gate is a tail statistic, and a context-shaped cell #: draws ~100x more elements than a 128-row one. Against the same float64 #: reference over three seeds at 16384 rows, the composition this replaces -#: reaches 2.91 row-max ulp (normed) and the op 2.94; at 8192 rows 1.89 / 1.89 -#: (`microbench/diag_band.py` in the candidate dir). Cells above 1024 rows are -#: therefore gated at 4; the RMS gate is unchanged. +#: reaches 2.91 row-max ulp (normed) and the op 2.94; at 8192 rows 1.89 / 1.89. +#: Cells above 1024 rows are therefore gated at 4; the RMS gate is unchanged. CONTEXT_ROWS, CONTEXT_ROW_ULP = 1024, 4.0 diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_block.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_block.py index 418852a9b93d..10c6bdc31fc8 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_block.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_block.py @@ -144,7 +144,7 @@ def _decode_state(m: int): buf[:m] = torch.tensor(_slots(m), dtype=torch.int32, device=DEV) si_all = buf[:m] qsl_long = torch.arange(m + 1, dtype=torch.int64, device=DEV) - return (si_all, qsl_long, False, m, 1, None, None), qsl_long + return (si_all, False, m, 1, None, None), qsl_long def _prefill_state(ctx_lens: tuple, n_dec: int, has_init: tuple): diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_gated_norm_pdl.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_gated_norm_pdl.py index c3718a602f32..7f45e90b392e 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_gated_norm_pdl.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_gated_norm_pdl.py @@ -2,9 +2,9 @@ # SPDX-License-Identifier: Apache-2.0 """GPU test for the gdn_decode_gated_norm_pdl entry. -Cells against the composition's fp64 reference (the recurrence entry's, then the -gated norm entry's) in the recurrence entry's band: torch's default for the bf16 -output and the fp32 pool, the pool compared whole so slots the call did not name are +Cells against the composition's fp64 reference (the gated delta rule's recurrence, +then the gated norm entry's) in torch's default band for the bf16 output and the fp32 +pool, the pool compared whole so slots the call did not name are gated too. Beyond the cells: agreement with the unfused composition it replaces (FlashInfer's recurrence, then the gated norm entry), PDL on and off bitwise, a stale-read check under graph replay (a producer rewrites the projection inside the @@ -13,13 +13,10 @@ import zlib +import flashinfer.gdn_decode import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.gated_delta_rule_decode_pretranspose import ( # noqa: E501 # noqa: E501 - _scattered, - gated_delta_rule_decode_pretranspose, -) 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, ) @@ -41,6 +38,13 @@ WIDTH = 2 * H * K + 2 * HV * V + 2 * HV +def _scattered(n: int, slots: int) -> tuple[int, ...]: + """`n` distinct slots of a `slots`-row pool, both ends first, then a stride-37 walk + over the interior (37 is coprime to every pool size here).""" + rest = tuple((i * 37 + 4) % (slots - 2) + 1 for i in range(max(n - 2, 0))) + return ((slots - 1, 0) + rest)[:n] + + def _build(spec, seed): g = torch.Generator(device=DEV).manual_seed(seed) rows = spec["rows"] @@ -94,7 +98,7 @@ def _unfused(proj, pool, slots, A_log, dt_bias, w, eps, h, k): rows = proj.shape[0] q, kk, v, z, a, b = torch.split(proj, [h * k, h * k, HV * V, HV * V, HV, HV], 1) o = torch.empty(rows, 1, HV, V, dtype=torch.bfloat16, device=DEV) - gated_delta_rule_decode_pretranspose( + flashinfer.gdn_decode.gated_delta_rule_decode_pretranspose( q=q.reshape(rows, 1, h, k), k=kk.reshape(rows, 1, h, k), v=v.reshape(rows, 1, HV, V), diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py index 2aba437bc7f9..b27964b8fd6e 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py @@ -2,9 +2,9 @@ # SPDX-License-Identifier: Apache-2.0 """GPU test for the gdn_decode_rows_gated_norm entry. -Cells against the composition's fp64 reference (the recurrence entry's, then the -gated norm entry's) in the recurrence entry's band: torch's default for the bf16 -output and the fp32 pool, the pool compared whole so slots the call did not name are +Cells against the composition's fp64 reference (the gated delta rule's recurrence, +then the gated norm entry's) in torch's default band for the bf16 output and the fp32 +pool, the pool compared whole so slots the call did not name are gated too. Beyond the cells: agreement with the unfused composition it replaces (FlashInfer's recurrence, then the gated norm entry) at decode batch sizes, PDL on and off bitwise, a stale-read check under graph replay (a producer rewrites the @@ -13,13 +13,10 @@ import zlib +import flashinfer.gdn_decode import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.gated_delta_rule_decode_pretranspose import ( # noqa: E501 - _scattered, - gated_delta_rule_decode_pretranspose, -) 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, ) @@ -44,6 +41,13 @@ WIDTH = 2 * H * K + 2 * HV * V + 2 * HV +def _scattered(n: int, slots: int) -> tuple[int, ...]: + """`n` distinct slots of a `slots`-row pool, both ends first, then a stride-37 walk + over the interior (37 is coprime to every pool size here).""" + rest = tuple((i * 37 + 4) % (slots - 2) + 1 for i in range(max(n - 2, 0))) + return ((slots - 1, 0) + rest)[:n] + + def _build(spec, seed): g = torch.Generator(device=DEV).manual_seed(seed) rows = spec["rows"] @@ -97,7 +101,7 @@ def _unfused(proj, pool, slots, A_log, dt_bias, w, eps, h, k): rows = proj.shape[0] q, kk, v, z, a, b = torch.split(proj, [h * k, h * k, HV * V, HV * V, HV, HV], 1) o = torch.empty(rows, 1, HV, V, dtype=torch.bfloat16, device=DEV) - gated_delta_rule_decode_pretranspose( + flashinfer.gdn_decode.gated_delta_rule_decode_pretranspose( q=q.reshape(rows, 1, h, k), k=kk.reshape(rows, 1, h, k), v=v.reshape(rows, 1, HV, V), diff --git a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_gated_delta_rule_decode_pretranspose.py b/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_gated_delta_rule_decode_pretranspose.py deleted file mode 100644 index 77bc3fd1732e..000000000000 --- a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_gated_delta_rule_decode_pretranspose.py +++ /dev/null @@ -1,167 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the gated_delta_rule_decode_pretranspose catalog entry (FlashInfer fp32-pool decode).""" - -import zlib - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.gated_delta_rule_decode_pretranspose import ( # noqa: E501 - QWEN38_GEOMETRY, - _scattered, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.gated_delta_rule_decode_pretranspose import ( # noqa: E501 - gated_delta_rule_decode_pretranspose as op, -) - -__extra_import_path__ = [".."] -from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name - -assert torch.cuda.is_available(), "gated_delta_rule_decode_pretranspose requires a CUDA device" - -DEV = "cuda" -_PAD = 64 - - -def _compare(out, ref) -> None: - """`(o, pool)` against the fp64 reference rounded once, at torch's default - dtype-aware band. Measured on the microbench inputs: o max abs 1.2e-4 at - |o| <= 6.2e-2 (identical to the Triton op's), state max abs 2.6e-7. The - pool is compared whole, so slots the call did not name are gated too.""" - (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.shape == ref_pool.shape and pool.dtype == ref_pool.dtype - torch.testing.assert_close(pool, ref_pool) - - -def _build(spec, seed): - """`__call__`'s arguments (by name) plus the buffers some layouts view.""" - g = torch.Generator(device=DEV).manual_seed(seed) - H, HV, K, V = spec["geometry"] - m, slots = spec["m"], spec["slots"] - mixed = torch.randn(m, 2 * H * K + HV * V, generator=g, device=DEV, dtype=torch.bfloat16) - qv, kv, vv = torch.split(mixed, [H * K, H * K, HV * V], 1) - q, k, v = ( - torch.reshape(qv, [m, 1, H, K]), - torch.reshape(kv, [m, 1, H, K]), - torch.reshape(vv, [m, 1, HV, V]), - ) - if spec.get("qkv", "views") == "contiguous": - q, k, v = q.contiguous(), k.contiguous(), v.contiguous() - a = torch.randn(m, HV, generator=g, device=DEV, dtype=torch.bfloat16).view(m, 1, HV) - b = torch.randn(m, HV, generator=g, device=DEV, dtype=torch.bfloat16).view(m, 1, HV) - A_log = torch.empty(HV, device=DEV).uniform_(-5.5, 1.0, generator=g) - dt_bias = torch.randn(HV, generator=g, device=DEV) * 4.0 - layout = spec.get("pool", "contiguous") - extras = {} - if layout == "padded": - backing = torch.randn(slots, HV * V * K + _PAD, generator=g, device=DEV) * 0.05 - pool = backing[:, : HV * V * K].view(slots, HV, V, K) - extras["backing"] = backing - elif layout == "layer_view": - cache = torch.randn(3, slots, HV, V, K, generator=g, device=DEV) * 0.05 - pool = cache[1] - extras["cache"] = cache - else: - pool = torch.randn(slots, HV, V, K, generator=g, device=DEV) * 0.05 - indices = spec.get("indices") or _scattered(m, slots) - indices = torch.tensor(indices, dtype=torch.int32, device=DEV) - output = None - if spec.get("output", "buffer") == "buffer": - output = torch.full((m, 1, HV, V), -7.0, device=DEV, dtype=torch.bfloat16) - args = dict( - q=q, - k=k, - v=v, - state=None, - A_log=A_log, - a=a, - dt_bias=dt_bias, - b=b, - scale=None, - output=output, - use_qk_l2norm=True, - initial_state=pool, - initial_state_indices=indices, - output_state_indices=None, - ) - 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) - args, extras = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) - pool = args["initial_state"] - pool_before = pool.clone() - cache_before = extras["cache"].clone() if "cache" in extras else None - slot_numel = pool[0].numel() - pad_before = extras["backing"][:, slot_numel:].clone() if "backing" in extras else None - ref = op.reference(**args) - - o, returned_pool = op(**args) - - assert returned_pool.data_ptr() == pool.data_ptr(), "the pool is advanced in place" - _compare((o, pool), ref) - if args["output"] is not None: - assert o.data_ptr() == args["output"].data_ptr(), "output= must be written in place" - named = {s for s in args["initial_state_indices"].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" - for slot in named: - assert not torch.equal(pool[slot], pool_before[slot]), slot - if cache_before is not None: - for other in (0, 2): - assert torch.equal(extras["cache"][other], cache_before[other]), other - if pad_before is not None: - assert torch.equal(extras["backing"][:, slot_numel:], pad_before), "slot padding written" - - -def test_repeated_calls_are_bitwise() -> None: - """Two identical calls from the same start give identical output and state.""" - results = [] - for _ in range(2): - args, _ = _build(dict(geometry=QWEN38_GEOMETRY, m=8, slots=40), seed=5) - o, pool = op(**args) - results.append((o.clone(), pool.clone())) - assert torch.equal(results[0][0], results[1][0]) and torch.equal(results[0][1], results[1][1]) - - -def test_matches_in_tree_triton_op() -> None: - """Same inputs through the in-tree Triton decode op this entry stands in - for: output and state agree at the default band (each is within it of fp64).""" - from tensorrt_llm._torch.modules.fla.fused_sigmoid_gating_recurrent import ( - fused_sigmoid_gating_delta_rule_update as triton_op, - ) - - H, HV, K, V = QWEN38_GEOMETRY - m = 32 - args, _ = _build(dict(geometry=QWEN38_GEOMETRY, m=m, slots=461), seed=11) - pool_t = args["initial_state"].clone() - cu = torch.arange(m + 1, device=DEV, dtype=torch.int64) - o_t = triton_op( - A_log=args["A_log"], - a=args["a"].view(m, HV), - dt_bias=args["dt_bias"], - softplus_beta=1.0, - softplus_threshold=20.0, - q=torch.reshape(args["q"], [1, m, H, K]), - k=torch.reshape(args["k"], [1, m, H, K]), - v=torch.reshape(args["v"], [1, m, HV, V]), - b=args["b"].view(m, HV), - initial_state_source=pool_t, - initial_state_indices=args["initial_state_indices"], - scale=None, - use_qk_l2norm_in_kernel=True, - cu_seqlens=cu, - output=None, - ) - o, pool = op(**args) - torch.testing.assert_close(o.view(1, m, HV, V), o_t) - torch.testing.assert_close(pool, pool_t) From 370bcd0116caf7a4d5310d1f566f7dd80d1eb972 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Mon, 19 Oct 2026 01:16:56 -0700 Subject: [PATCH 11/28] [TRTLLM-17033][test] Run most qwen4_exp accuracy gates post-merge Every gate builds the full Qwen3.8-Flash-Next-NVFP4 engine (a gsm8k gate is roughly 8-13 min on one GB300), so the 26 pre-merge gates cost several GPU hours per CI run. Pre-merge keeps main's three gates (gsm8k; gsm8k with one MTP draft and the acceptance-length floor; mmmu with MTP) and adds the one gate that runs the sparse attention path: 3k-6k-token pass keys with three drafts and IndexShare over bf16 KV. Post-merge runs one gate per opt-in variant: MTP3 alone and over CuteDSL experts, MTP7 over long pass keys, FP8 and NVFP4 KV over the pass keys, the bf16 recurrent state at MTP3, MTP7 all rejected by forced acceptance, video, and the breakable prefill graph against eager. The other gates stay in test_qwen4_exp.py for on-demand runs. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../test_lists/test-db/l0_b300.yml | 64 ++++++++----------- 1 file changed, 25 insertions(+), 39 deletions(-) diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index d73cb66fc3a9..4d0e2f9dc07a 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -135,50 +135,19 @@ l0_b300: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_chunked_prefill[quant_dtype=fp8-kv_cache_reuse=True-fp8kv=True-overlap_scheduler=True] # Cover nvbugs 6084445 - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer[enable_chunked_prefill=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer[enable_chunked_prefill=True] - # TIMEOUT measured: 2259 cases in 11m55s on one GB300 (the cpu_only files run on the CPU - # stage instead), so 30 is 2.5x headroom. + # TIMEOUT measured: 3122 cases in 13m54s on one GB300 (the cpu_only files run on the CPU + # stage instead), so 30 is 2.2x headroom. - unittest/_torch/modeling_v2 --ignore=unittest/_torch/modeling_v2/comm TIMEOUT (30) - accuracy/modeling_v2/test_gpt_oss.py::TestModelingV2GptOss120bSm10xDefault::test_gsm8k[tp1] - # qwen4_exp target: gsm8k, and gsm8k with the trunk's routed experts on CuteDSL; - # gsm8k with MTP plus the acceptance-length floor, the only gate that sees a - # miscomputed draft layer, at one draft over bf16 and fp8 KV, at three drafts - # alone and over CuteDSL experts, and at five and seven drafts (the verify step's - # six and eight rows per request); 3k-6k-token pass key prompts, the only cases that - # run the sparse attention path (IndexShare over bf16 KV, fp8 KV), and 3k-12k-token ones - # at seven drafts; seven drafts all rejected and all accepted by forced acceptance at - # batch 32 over long prompts; the recurrent-state verify path that - # TRTLLM_USE_GDN_REPLAY=0 selects; the recurrent state stored in bf16 without speculation, - # at three drafts and at seven drafts over long prompts; mmmu with MTP and synthetic - # videos, the only gates that reach the vision tower; and the opt-in breakable prefill - # CUDA graph: gsm8k, gsm8k with three drafts and the length floor, 3k-6k-token pass keys - # prefilled in replayed 2048-token chunks, mmmu, and replayed against eager greedy - # continuations at the bucket edges and with context and generation rows in one batch. + # qwen4_exp target, pre-merge: gsm8k; gsm8k with one MTP draft plus the + # acceptance-length floor, the only gate that sees a miscomputed draft layer; + # mmmu with MTP, which reaches the vision tower; and 3k-6k-token pass key prompts + # with three drafts and IndexShare over bf16 KV, the only prompts long enough to run + # the sparse attention path. The opt-in variants run post-merge (below). - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_gsm8k - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_gsm8k_cutedsl_moe - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp1_gsm8k_and_acceptance - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp1_fp8_kv_gsm8k_and_acceptance - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp1_nvfp4_kv_gsm8k_and_acceptance - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_gsm8k_and_acceptance - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_cutedsl_gsm8k_and_acceptance - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp5_gsm8k_and_acceptance - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp7_gsm8k_and_acceptance - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_long_context_passkey[index_share] - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_long_context_passkey[fp8_kv] - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_long_context_passkey[nvfp4_kv] - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp7_long_context_passkey - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp7_forced_acceptance_extremes[all_rejected] - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp7_forced_acceptance_extremes[all_accepted] - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_gsm8k_bf16_ssm_state - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_bf16_ssm_state_gsm8k_and_acceptance - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp7_bf16_ssm_state_long_context_passkey - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp1_gdn_legacy_verify - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp1_mmmu - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp1_video - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_gsm8k_breakable_prefill_cuda_graph - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_breakable_prefill_gsm8k_and_acceptance - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_long_context_passkey[breakable_prefill] - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp1_mmmu_breakable_prefill - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_breakable_prefill_matches_eager + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_long_context_passkey[index_share] - condition: ranges: system_gpu_count: @@ -199,3 +168,20 @@ l0_b300: - unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k "TRTLLM" - unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k "CUTEDSL" - unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k "DEEPGEMM" + # qwen4_exp target, post-merge: one gate per opt-in variant -- three drafts, alone and + # over CuteDSL experts; seven drafts over 3k-12k-token pass keys (the verify step's + # eight rows per request); fp8 and NVFP4 KV over the pass keys (their sparse attention + # path); the recurrent state in bf16 at three drafts; seven drafts all rejected by + # forced acceptance (the speculative rollback end to end); synthetic videos; and the + # breakable prefill CUDA graph replayed against eager greedy continuations at the bucket + # edges and with context and generation rows in one batch. The other gates in + # test_qwen4_exp.py run on demand. + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_gsm8k_and_acceptance + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_cutedsl_gsm8k_and_acceptance + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp7_long_context_passkey + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_long_context_passkey[fp8_kv] + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_long_context_passkey[nvfp4_kv] + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_bf16_ssm_state_gsm8k_and_acceptance + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp7_forced_acceptance_extremes[all_rejected] + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp1_video + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_breakable_prefill_matches_eager From 1f515cf0b0f344ce8b1b32ae3a5984096b14a259 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Mon, 19 Oct 2026 19:36:36 -0700 Subject: [PATCH 12/28] [TRTLLM-17033][fix] Give the beam-search admission stub the model engine the model hook reads PyExecutor._validate_request now also calls the model's optional validate_request, which it finds through self.model_engine.model. The executor stub of test_vbws_rejects_decreasing_beam_width_array listed every member _validate_request read before and failed with AttributeError on model_engine; it now carries a model engine without a model. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- tests/unittest/_torch/sampler/test_beam_search.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index 43a80f733f87..ea9a23ef8f68 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -2014,7 +2014,8 @@ def test_vbws_rejects_decreasing_beam_width_array(beam_width_array: list[int], assert request.py_beam_width == max(beam_width_array) # Exercise the real beam admission checks with unrelated cache, transport, - # token-range, sampler, and budget validation disabled in the executor stub. + # token-range, sampler, model, and budget validation disabled in the + # executor stub. executor = types.SimpleNamespace( _is_kv_manager_v2=True, max_beam_width=request.py_beam_width, @@ -2023,6 +2024,7 @@ def test_vbws_rejects_decreasing_beam_width_array(beam_width_array: list[int], _configure_csa2_remote_tail=lambda _request: None, _validate_token_id_range=lambda _request: None, sampler=types.SimpleNamespace(validate_request=lambda _request: None), + model_engine=types.SimpleNamespace(model=None), _validate_request_budget=lambda _request: None, ) validate = functools.partial( From 7c9e76775f49a46673e46b7aa6fe32e8d6379f37 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Mon, 19 Oct 2026 19:44:42 -0700 Subject: [PATCH 13/28] [TRTLLM-16460][fix] Qwen3.8-Flash-Next target: keep MTP verify steps on the composed attention prologue qwen4_exp_attn_decode_prologue requires one row per request: at a compression group's last token it pools the group's earlier raw keys, and the first token's coordinates, from the index plane before its PDL wait, as earlier steps wrote them. A verify step packs draft_len + 1 rows of a request into one launch, so a group ending on a later row read the earlier rows' slots before the same launch wrote them. The index scores' PDL prefetch has the same gap: with one row per request the prologue writes no compressed key but the one the entry re-reads after its wait; in a verify group an earlier row can complete a group a later row scores. Verify steps (stage) now run the composed prologue, whose index store and compression are separate launches and which stages the index plane for the rollback itself; use_pdl and prefetch follow fused_prologue, so they are off there too. Steps with one row per request are unchanged. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../modeling.py | 20 ++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py index 4f7b5f2e4a3c..803e6b30f9f0 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py @@ -4677,9 +4677,12 @@ def _attention_block( advanced, so a request that grows past the threshold finds its prefix already compressed; above it the dense core is replaced outright and the main K/V append becomes this layer's own call, because the op that would - have done it is not the one running. Above the threshold over a bf16 pool - everything up to the index scores is one launch, - `qwen4_exp_attn_decode_prologue`.""" + have done it is not the one running. Above the threshold over a bf16 pool, + on a step with one row per request, everything up to the index scores is one + launch, `qwen4_exp_attn_decode_prologue`. A verify step stays on the composed + prologue: at a group's last token that launch reads the group's earlier keys + from the index plane, and in a verify group the rows before it are written by + the same launch.""" core = self.core ( req_idx, @@ -4716,13 +4719,11 @@ def _attention_block( index_section = (_QSA_MROPE_SECTION or (0, 0, 0)) if _QSA_INDEX_USE_MROPE else (0, 0, 0) index_cache = core._bound.index_cache[i] kv = core._bound.kv[i] - fused_prologue = sparse and kv.k_scale is None and kv.k.dtype == proj.dtype + fused_prologue = sparse and not stage and kv.k_scale is None and kv.k.dtype == proj.dtype if fused_prologue: # The norm + rope of both planes, the K/V append the dense op would have # owned, and the index branch's raw-key and coordinate store and group # compression in one launch on this stream. - if stage: - core.spec_rollback.stage_index(index_cache, core._bound.position_cache) qwen4_exp_attn_decode_prologue( proj, q_index, @@ -4782,9 +4783,10 @@ def _attention_block( only_visible_blocks=True, context_rows=False, # Its predecessor is the prologue launch and its consumer the top-k, - # which waits before reading the scores. The prologue writes no - # compressed key but the one the entry re-reads after its wait, so at - # decode-sized batches the pages and keys are read before it. + # which waits before reading the scores. With one row per request the + # prologue writes no compressed key but the one the entry re-reads + # after its wait, so at decode-sized batches the pages and keys are + # read before it. use_pdl=fused_prologue, prefetch=fused_prologue and num_tokens <= _SPARSE_DECODE_MAX_ROWS, ) From bf7b60fd9bcf1b89f4521942fbe4eb0f2f507f2e Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Mon, 19 Oct 2026 19:44:53 -0700 Subject: [PATCH 14/28] [TRTLLM-16460][fix] Qwen3.8-Flash-Next target: gated norm sends a z that does not cover x to the checked op From 32768 rows rms_norm_gated_token_major runs its own multi-row launch, which skips torch.ops.trtllm.rms_norm_gated_token_major and with it the assertion that z matches x. Row r reads z[r // heads, r % heads], so a z with fewer tokens was read past its end. The launch now also requires x.shape[0] == z.shape[0] * z.shape[1] and z.shape[2] == N; anything else goes to the registered op, which rejects it as the op's note says. The new test drives a 683-token x with a z one token short. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../ops/norm/rms_norm_gated_token_major.py | 3 +++ .../norm/test_modeling_v2_rms_norm_gated_token_major.py | 9 +++++++++ 2 files changed, 12 insertions(+) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/rms_norm_gated_token_major.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/rms_norm_gated_token_major.py index a9a667c28e0d..0cce6669186d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/rms_norm_gated_token_major.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/rms_norm_gated_token_major.py @@ -29,9 +29,12 @@ def _prefill_multirow(x: torch.Tensor, z: torch.Tensor, gate_activation: str) -> bool: + # A `z` that does not cover `x` goes to the registered op, which rejects it. n = x.shape[-1] return ( x.shape[0] >= _PREFILL_MIN_ROWS + and x.shape[0] == z.shape[0] * z.shape[1] + and z.shape[2] == n and x.stride(-1) == 1 and z.stride(2) == 1 and z.stride(1) == n diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_rms_norm_gated_token_major.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_rms_norm_gated_token_major.py index f3a733751962..5d1725380046 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_rms_norm_gated_token_major.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_rms_norm_gated_token_major.py @@ -177,3 +177,12 @@ def test_rejects_z_that_does_not_cover_x() -> None: 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) + + +def test_prefill_launch_rejects_z_that_does_not_cover_x() -> None: + """From 32768 rows the op runs its own multi-row launch; a `z` one token short + still reaches the registered op's check instead of being read past its end.""" + x, z, weight, eps, _, act = _build(dict(op.CELLS[1].spec, tokens=683), seed=12) + assert x.shape[0] >= 32768, x.shape + with pytest.raises(AssertionError, match="does not match x shape"): + op(x, z[:-1], weight, eps, None, act) From 66ab699a5600c76bedc4459569cbcc64affbc828 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Mon, 19 Oct 2026 23:01:26 -0700 Subject: [PATCH 15/28] [TRTLLM-17033][chore] modeling_v2 routing: give each family's route the ModelConfig and LLM args The router index derived the engine choices only the qwen4_exp target routes on (EngineFacts: model dtype, MoE backend, recurrent-state dtype, GDN replay, prefill CUDA graph backend) and passed them to every family. It now passes the ModelConfig and the TorchLlmArgs themselves, as model_config= and llm_args=, beside the kv_cache_quant_algo argument every route already takes, and each family derives what it routes on: - qwen4_exp/routing.py owns EngineChoices and engine_choices(), the former EngineFacts and engine_facts(), and declines from them as before; its KV check reads kv_cache_quant_algo, so an fp8_ds_mla cache is refused by name; - gpt_oss and deepseek_v3 take both arguments, typed instead of Any, and read neither; - the loader's effective_kv_cache_quant, which only EngineFacts called, is back inline. test_qwen4_exp_tp1_matches also pins TRTLLM_USE_GDN_REPLAY, since its mtp7 case needs the cached replay the variable can turn off. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../_experimental/modeling_v2/README.md | 11 ++-- .../modeling_v2/_router_index.py | 51 +++---------------- .../modeling_v2/models/deepseek_v3/routing.py | 12 +++-- .../modeling_v2/models/gpt_oss/routing.py | 12 +++-- .../modeling_v2/models/qwen4_exp/routing.py | 49 ++++++++++++++++-- .../_torch/pyexecutor/model_loader.py | 14 ++--- .../modeling_v2/test_modeling_v2_routing.py | 13 +++-- 7 files changed, 88 insertions(+), 74 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md index 8528a8c20827..4e86b1664588 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md @@ -37,9 +37,11 @@ exists to catch. The decision is one function, `route`, in the family's `routing.py`. It reads the facts that fix which code runs -- the checkpoint's config shape, the parallel mapping, the GPU's compute capability, the speculative-decoding -config and the engine choices in `EngineFacts` (KV cache quantization, MoE -backend, recurrent-state cache dtype, prefill CUDA graph backend) -- and either returns a target class or -raises `NoTarget` with the reason in words. A target declines what its +config, the KV cache algorithm the engine will run with, and the engine choices +it derives from the `ModelConfig` and the LLM arguments it is also given +(qwen4_exp: MoE backend, recurrent-state cache dtype, prefill CUDA graph +backend) -- and either returns a target class or raises `NoTarget` with the +reason in words. A target declines what its constructor would refuse, so that `auto` falls back instead of failing mid-build. Tuning knobs such as `max_batch_size`, `max_num_tokens` and the CUDA graph bucket lists are not routing inputs: they do not change which code runs, so a @@ -70,7 +72,8 @@ LLM(model=..., modeling_v2=...) -> ModelLoader | _router_index.py +-- architectures[0] -> routing module models//routing.py +-- route(config, mapping, compute_capability, - | spec_config, kv_cache_quant_algo, engine) + | spec_config, kv_cache_quant_algo, + | model_config, llm_args) | -> target class, or NoTarget | config.modeling_v2_target = class diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py index c9b4c681f79c..cd2a1169d7bd 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py @@ -18,10 +18,11 @@ * The routing module's ``route`` owns the rest of the decision as one forward-reading tree over the facts that fix which code runs: the checkpoint's config, the parallel mapping, the GPU's compute capability, the - speculative-decoding config, the KV-cache algorithm the engine will run with - and the engine choices in :class:`EngineFacts`. It returns the target class, - or raises :class:`NoTarget` with the reason in words. Reading that one file - tells you where any configuration lands. + speculative-decoding config and the KV-cache algorithm the engine will run + with. It also receives the ``ModelConfig`` and the LLM arguments, from which a + family derives any other engine choice its targets depend on. It returns the + target class, or raises :class:`NoTarget` with the reason in words. Reading + that one file tells you where any configuration lands. Tuning knobs -- ``max_batch_size``, ``max_num_tokens``, CUDA graph settings -- are deliberately not routing inputs. They do not change which code runs, so a @@ -38,7 +39,7 @@ import enum import os from importlib import import_module -from typing import TYPE_CHECKING, Any, NamedTuple, Optional, Tuple +from typing import TYPE_CHECKING, Optional, Tuple from tensorrt_llm.logger import logger @@ -166,43 +167,6 @@ def _kv_cache_quant_algo(config: "ModelConfig", llm_args: "TorchLlmArgs"): return {"fp8": QuantAlgo.FP8, "nvfp4": QuantAlgo.NVFP4}.get(dtype, dtype) -class EngineFacts(NamedTuple): - """Engine choices that fix which code runs but are not in the checkpoint's config. - - A target's constructor refuses what it cannot run, but by then ``route`` has - already claimed the deployment and ``auto`` can no longer fall back; a route - that reads these declines first. - """ - - torch_dtype: Any # torch.dtype the engine runs the model in - kv_cache_quant_algo: Any # QuantAlgo or None: the checkpoint's under "auto", else the override - moe_backend: str # `moe_config.backend` as requested, before AUTO resolves - mamba_ssm_cache_dtype: Any # torch.dtype of the recurrent-state cache - gdn_replay: bool # TRTLLM_USE_GDN_REPLAY, which the engine reads for eligible MTP workloads - prefill_cuda_graph_backend: str # `prefill_cuda_graph_backend`: disabled, piecewise, breakable - - -def engine_facts(config: "ModelConfig", llm_args: "TorchLlmArgs") -> EngineFacts: - # Imported here: the model loader imports this module. - from tensorrt_llm._torch.pyexecutor.model_loader import ( - effective_kv_cache_quant, - effective_mamba_ssm_cache_dtype, - ) - from tensorrt_llm._torch.utils import is_gdn_replay_enabled - - kv_cache_config = llm_args.kv_cache_config - return EngineFacts( - torch_dtype=config.torch_dtype, - kv_cache_quant_algo=effective_kv_cache_quant(config, kv_cache_config.dtype), - moe_backend=config.moe_backend_requested, - mamba_ssm_cache_dtype=effective_mamba_ssm_cache_dtype( - config, kv_cache_config.mamba_ssm_cache_dtype - ), - gdn_replay=is_gdn_replay_enabled(), - prefill_cuda_graph_backend=str(llm_args.prefill_cuda_graph_backend.value), - ) - - def modeling_v2_resolve(config: "ModelConfig", llm_args: "TorchLlmArgs") -> Optional[type]: """Decide which modeling_v2 target, if any, this deployment builds. @@ -233,7 +197,8 @@ def modeling_v2_resolve(config: "ModelConfig", llm_args: "TorchLlmArgs") -> Opti _device_compute_capability(), config.spec_config, _kv_cache_quant_algo(config, llm_args), - engine=engine_facts(config, llm_args), + model_config=config, + llm_args=llm_args, ) except NoTarget as why: if mode is ModelingV2Mode.REQUIRE: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/routing.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/routing.py index 8d0502f23bcc..6f4dd9fd6c2f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/routing.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/routing.py @@ -11,7 +11,7 @@ from __future__ import annotations -from typing import Any, Optional, Tuple +from typing import TYPE_CHECKING, Any, Optional, Tuple from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig from tensorrt_llm.quantization.mode import QuantAlgo @@ -69,6 +69,10 @@ } ) +if TYPE_CHECKING: + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + def route( pretrained_config: Any, @@ -76,7 +80,8 @@ def route( compute_capability: Tuple[int, int], spec_config: Optional[Any], kv_cache_quant_algo: Optional[Any], - engine: Optional[Any] = None, + model_config: Optional[ModelConfig] = None, + llm_args: Optional[TorchLlmArgs] = None, ) -> type: """Return the target class for this deployment, or raise ``NoTarget`` saying why none fits. @@ -84,7 +89,8 @@ def route( (see ``_router_index._kv_cache_quant_algo``), not the checkpoint's alone. Each target is imported where it is returned, so declining costs no import - of modeling code. + of modeling code. No engine choice in ``model_config`` or ``llm_args`` changes + which of these targets runs, so neither is read. """ # The checkpoint is recognized by its config shape; see the gpt_oss routing # module for what shape-sniffing does and does not pin. q_lora_rank is in diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/routing.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/routing.py index e18cc12e1644..7bd51d78e37a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/routing.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/routing.py @@ -10,7 +10,7 @@ from __future__ import annotations -from typing import Any, Optional, Tuple +from typing import TYPE_CHECKING, Any, Optional, Tuple from ..._router_index import NoTarget, tp_reaches_across_nodes @@ -42,6 +42,10 @@ } ) +if TYPE_CHECKING: + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + def route( pretrained_config: Any, @@ -49,7 +53,8 @@ def route( compute_capability: Tuple[int, int], spec_config: Optional[Any], kv_cache_quant_algo: Optional[Any], - engine: Optional[Any] = None, + model_config: Optional[ModelConfig] = None, + llm_args: Optional[TorchLlmArgs] = None, ) -> type: """Return the target class for this deployment, or raise ``NoTarget`` saying why none fits. @@ -57,7 +62,8 @@ def route( (see ``_router_index._kv_cache_quant_algo``), not the checkpoint's alone. Each target is imported where it is returned, so declining costs no import - of modeling code. + of modeling code. No engine choice in ``model_config`` or ``llm_args`` changes + which of these targets runs, so neither is read. """ # The checkpoint is recognized by its config shape, the upstream idiom # (``is_mla``, ``is_nemotron_hybrid`` do the same). A fine-tune with the diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/routing.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/routing.py index 2cecbbc5fad0..ef2f31f01234 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/routing.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/routing.py @@ -12,14 +12,18 @@ from __future__ import annotations -from typing import Any, Optional, Tuple +from typing import TYPE_CHECKING, Any, NamedTuple, Optional, Tuple import torch from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig from tensorrt_llm.quantization.mode import QuantAlgo -from ..._router_index import EngineFacts, NoTarget +from ..._router_index import NoTarget + +if TYPE_CHECKING: + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs # (layers, hidden, experts, attention heads, linear value heads, vocab). Shape-sniffing # also matches fine-tunes; the accuracy gate pins identity. @@ -44,8 +48,36 @@ _SSM_STATE_DTYPES = (torch.float32, torch.bfloat16) +class EngineChoices(NamedTuple): + """Engine choices that fix which code the target runs but are not in the checkpoint's + config, as the engine will resolve them.""" + + torch_dtype: torch.dtype # the dtype the engine runs the model in + moe_backend: str # `moe_config.backend` as requested, before AUTO resolves + mamba_ssm_cache_dtype: torch.dtype # the recurrent-state cache's dtype + gdn_replay: bool # TRTLLM_USE_GDN_REPLAY, which the engine reads for eligible MTP workloads + prefill_cuda_graph_backend: str # `prefill_cuda_graph_backend`: disabled, piecewise, breakable + + +def engine_choices(model_config: ModelConfig, llm_args: TorchLlmArgs) -> EngineChoices: + # Imported on the claimed branch only, as the target is. + from tensorrt_llm._torch.pyexecutor.model_loader import effective_mamba_ssm_cache_dtype + from tensorrt_llm._torch.utils import is_gdn_replay_enabled + + kv_cache_config = llm_args.kv_cache_config + return EngineChoices( + torch_dtype=model_config.torch_dtype, + moe_backend=model_config.moe_backend_requested, + mamba_ssm_cache_dtype=effective_mamba_ssm_cache_dtype( + model_config, kv_cache_config.mamba_ssm_cache_dtype + ), + gdn_replay=is_gdn_replay_enabled(), + prefill_cuda_graph_backend=str(llm_args.prefill_cuda_graph_backend.value), + ) + + def _check_engine( - spec_config: Optional[Any], kv_cache_quant_algo: Optional[Any], engine: Optional[EngineFacts] + spec_config: Optional[Any], kv_cache_quant_algo: Optional[Any], engine: Optional[EngineChoices] ) -> None: """Decline what the target's constructor or first forward would refuse. @@ -108,7 +140,8 @@ def route( compute_capability: Tuple[int, int], spec_config: Optional[Any], kv_cache_quant_algo: Optional[Any], - engine: Optional[EngineFacts] = None, + model_config: Optional[ModelConfig] = None, + llm_args: Optional[TorchLlmArgs] = None, ) -> type: """Return the target class for this deployment, or raise ``NoTarget`` saying why none fits. @@ -116,7 +149,8 @@ def route( (see ``_router_index._kv_cache_quant_algo``), not the checkpoint's alone. The target is imported where it is returned, so declining costs no import - of modeling code. + of modeling code. With ``model_config`` and ``llm_args`` it also declines the + engine choices (`EngineChoices`) the target's constructor would refuse. """ text = getattr(pretrained_config, "text_config", None) vision = getattr(pretrained_config, "vision_config", None) @@ -179,6 +213,11 @@ def route( # isort: on # fmt: on + engine = ( + engine_choices(model_config, llm_args) + if model_config is not None and llm_args is not None + else None + ) _check_engine(spec_config, kv_cache_quant_algo, engine) return ModelingV2Qwen38FlashNextNvfp4Sm103Tp1 diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 6cecef79e926..c53cb02cd55e 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -158,15 +158,6 @@ def validate_and_set_mamba_ssm_cache_dtype( config.quant_config.mamba_ssm_philox_rounds = mamba_ssm_philox_rounds -def effective_kv_cache_quant(model_config: ModelConfig, - pyt_kv_cache_dtype: str): - """The KV cache quantization the engine runs with: the checkpoint's under - "auto", otherwise what the explicit dtype maps to.""" - if pyt_kv_cache_dtype == "auto": - return model_config.quant_config.kv_cache_quant_algo - return _KV_CACHE_MAP.get(pyt_kv_cache_dtype, None) - - def validate_and_set_kv_cache_quant(model_config: ModelConfig, pyt_kv_cache_dtype: str) -> None: logger.info( @@ -178,8 +169,9 @@ def validate_and_set_kv_cache_quant(model_config: ModelConfig, valid_pyt_quant = bool(pyt_kv_cache_dtype in _VALID_KV_CACHE_DTYPES) mapped_pyt_quant = _KV_CACHE_MAP.get(pyt_kv_cache_dtype, None) - effective_quant = effective_kv_cache_quant(model_config, pyt_kv_cache_dtype) - if (effective_quant in (QuantAlgo.NVFP4, QuantAlgo.NVFP4.value) + effective_kv_cache_quant = (kv_cache_quant if pyt_kv_cache_dtype == "auto" + else mapped_pyt_quant) + if (effective_kv_cache_quant in (QuantAlgo.NVFP4, QuantAlgo.NVFP4.value) and not supports_fp4_mla_attention(model_config) and torch.cuda.is_available() and get_sm_version() == 107): logger.warning( 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 0da317f4c8ec..b47b5836a1f8 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py @@ -33,7 +33,6 @@ from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( NoTarget, - engine_facts, modeling_v2_resolve, ) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp import ( @@ -564,7 +563,9 @@ def _qwen4_exp_sm(_on_sm10x): ], ids=["no-spec", "mtp1", "mtp3", "mtp7"], ) -def test_qwen4_exp_tp1_matches(spec_config): +def test_qwen4_exp_tp1_matches(monkeypatch, spec_config): + # mtp7 needs the GDN cached replay, which TRTLLM_USE_GDN_REPLAY=0 turns off. + monkeypatch.delenv("TRTLLM_USE_GDN_REPLAY", raising=False) config = _model_config(_qwen4_exp_config(), spec_config=spec_config) # The recurrent-state cache at either dtype the target serves; `auto` is bf16 here. for dtype in ("float32", "bfloat16", "auto"): @@ -655,9 +656,11 @@ def test_qwen4_exp_quantized_kv_cache_dtypes_match(algo): @pytest.mark.usefixtures("_qwen4_exp_sm") @pytest.mark.parametrize("algo", [QuantAlgo.INT8, QuantAlgo.FP8_BLOCK_SCALES]) def test_qwen4_exp_other_kv_cache_algorithms_are_declined(algo): - facts = engine_facts(_model_config(_qwen4_exp_config()), _llm_args("auto")) - with pytest.raises(NoTarget, match="kv_cache_quant_algo"): - qwen4_exp_routing._check_engine(None, facts._replace(kv_cache_quant_algo=algo)) + choices = qwen4_exp_routing.engine_choices( + _model_config(_qwen4_exp_config()), _llm_args("auto") + ) + with pytest.raises(NoTarget, match="KV cache"): + qwen4_exp_routing._check_engine(None, algo, choices) @pytest.mark.usefixtures("_qwen4_exp_sm") From 7c8b176b576649e8673e814f9aec136a960de259 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Mon, 19 Oct 2026 19:45:58 -0700 Subject: [PATCH 16/28] [TRTLLM-16460][chore] Qwen3.8-Flash-Next target: own low_m_bf16_mm and flashinfer_chunk_gated_delta_rule Both entries have this target as their only caller, and low_m_bf16_mm's kernel table is tuned for its shapes. The README's ownership rule puts an entry one target calls under that target's ops/, so they move from catalog/gemm and catalog/ssm to the target's ops/gemm and ops/ssm, their tests to the target's test tree; catalog/ssm, which only held the GDN entry, goes. No code changes beyond the import paths. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../_experimental/modeling_v2/catalog/ssm/__init__.py | 3 --- .../qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py | 8 +++----- .../ops/gemm/bf16_gemv_pdl.py | 2 +- .../ops}/gemm/low_m_bf16_mm.py | 8 ++++++-- .../ops}/ssm/flashinfer_chunk_gated_delta_rule.py | 3 +-- .../gemm/test_modeling_v2_low_m_bf16_mm.py | 10 ++++++---- ...st_modeling_v2_flashinfer_chunk_gated_delta_rule.py | 8 ++++---- .../ssm/test_modeling_v2_gdn_in_proj_conv1d_pdl.py | 4 +++- 8 files changed, 24 insertions(+), 22 deletions(-) delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/__init__.py rename tensorrt_llm/_torch/_experimental/modeling_v2/{catalog => models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops}/gemm/low_m_bf16_mm.py (99%) rename tensorrt_llm/_torch/_experimental/modeling_v2/{catalog => models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops}/ssm/flashinfer_chunk_gated_delta_rule.py (99%) rename tests/unittest/_torch/modeling_v2/{ => models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1}/gemm/test_modeling_v2_low_m_bf16_mm.py (94%) rename tests/unittest/_torch/modeling_v2/{ => models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1}/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py (93%) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/__init__.py deleted file mode 100644 index bbf1df41d716..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""SSM entries.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py index 803e6b30f9f0..49c25e221fa4 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py @@ -46,7 +46,7 @@ layer's FP8 block-scale experts run on trtllm-gen under either. Under ``TRTLLM`` the prefill-sized launches (``_CUTE_DSL_MOE_MIN_ROWS``) run the CuteDSL pipeline over a second copy of the experts in its layout. A graph-captured decode step routes its block - projections through ``catalog/gemm/low_m_bf16_mm`` (FlashInfer low-M kernels at + projections through ``ops/gemm/low_m_bf16_mm`` (FlashInfer low-M kernels at M <= 32 on the shapes it lists, cuBLAS otherwise); see ``DecodeTarget.forward``. Integration facts this file depends on: @@ -124,13 +124,9 @@ thop_attention, ) from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm -from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.low_m_bf16_mm import low_m_bf16_mm from tensorrt_llm._torch._experimental.modeling_v2.catalog.graph.eager_on_graph import ( eager_on_graph, ) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.flashinfer_chunk_gated_delta_rule import ( # noqa: E501 - flashinfer_chunk_gated_delta_rule, -) 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 @@ -170,6 +166,7 @@ from .ops.embedding.ngram_hash import ngram_hash from .ops.gemm.bf16_gemv_pdl import TILING as _GEMV_TILED_K from .ops.gemm.bf16_gemv_pdl import bf16_gemv_pdl +from .ops.gemm.low_m_bf16_mm import low_m_bf16_mm from .ops.hc.qwen4_exp_gated_residual import qwen4_exp_gated_residual from .ops.hc.qwen4_exp_hc_decode_first_cute import qwen4_exp_hc_decode_first_cute from .ops.hc.qwen4_exp_hc_decode_mix_fp4_cute import ( @@ -211,6 +208,7 @@ from .ops.ssm.cu_seqlens_from_seq_lens import cu_seqlens_from_seq_lens from .ops.ssm.dilated_short_conv_paged import dilated_short_conv_paged from .ops.ssm.extract_transpose_prefill_slice import extract_transpose_prefill_slice +from .ops.ssm.flashinfer_chunk_gated_delta_rule import flashinfer_chunk_gated_delta_rule from .ops.ssm.fused_gdn_gating import fused_gdn_gating from .ops.ssm.fused_gdn_post_conv import fused_gdn_post_conv from .ops.ssm.fused_recurrent_gated_delta_rule_cached_replay_update import ( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/bf16_gemv_pdl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/bf16_gemv_pdl.py index 68163792965f..82ab484714f2 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/bf16_gemv_pdl.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/bf16_gemv_pdl.py @@ -311,7 +311,7 @@ def _cell(why: str, m: int, n: int, k: int, pdl: bool = True) -> Cell: class Bf16GemvPdl(OpWrapper): """`mat_a @ mat_b` in bf16 with fp32 accumulation, one PDL launch per call. - A drop-in for `catalog/gemm/low_m_bf16_mm` at decode row counts: `mat_a` is + A drop-in for `ops/gemm/low_m_bf16_mm` at decode row counts: `mat_a` is dense row-major `[M, K]` with `M <= MAX_ROWS`, `mat_b` the dense column-major `[K, N]` view of a stored `[N, K]` weight (its `.t()`), the result a fresh row-major `[M, N]` bf16 tensor. With `pdl` (and PDL available) the kernel is diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/low_m_bf16_mm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/low_m_bf16_mm.py similarity index 99% rename from tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/low_m_bf16_mm.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/low_m_bf16_mm.py index 238af508a3cd..133d99940b7e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/low_m_bf16_mm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/low_m_bf16_mm.py @@ -27,8 +27,12 @@ from flashinfer.gemm.kernels import dense_bf16_gemm_sm100_splitk as _fi_splitk import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* - -from .._op import Arch, Cell, OpWrapper, true_fp32_matmul +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) class _Plan(NamedTuple): diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/flashinfer_chunk_gated_delta_rule.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/flashinfer_chunk_gated_delta_rule.py similarity index 99% rename from tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/flashinfer_chunk_gated_delta_rule.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/flashinfer_chunk_gated_delta_rule.py index f4d3e1e69849..8a64fa713223 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/flashinfer_chunk_gated_delta_rule.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/flashinfer_chunk_gated_delta_rule.py @@ -21,12 +21,11 @@ import torch +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper from tensorrt_llm._torch.modules.fla.flashinfer_chunk import ( chunk_gated_delta_rule as _fi_chunk_gated_delta_rule, ) -from .._op import Arch, Cell, OpWrapper - #: qwen3.8-flash-next-nvfp4's Gated DeltaNet at tp1: (H, HV, K, V) -- 16 q/k #: heads, 48 v heads, head dims 128. QWEN38_GEOMETRY = (16, 48, 128, 128) diff --git a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_low_m_bf16_mm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_low_m_bf16_mm.py similarity index 94% rename from tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_low_m_bf16_mm.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_low_m_bf16_mm.py index c0b4dfd69ac4..f788ef2dffcb 100644 --- a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_low_m_bf16_mm.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_low_m_bf16_mm.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the low_m_bf16_mm catalog entry. +"""GPU test for the low_m_bf16_mm entry. Every cell names the branch `plan` must pick, so a table edit that moves a shape to another kernel fails here before it can change a number silently. @@ -14,12 +14,14 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.low_m_bf16_mm import ( +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.gemm.low_m_bf16_mm import ( # noqa: E501 low_m_bf16_mm as op, ) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.low_m_bf16_mm import plan +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.gemm.low_m_bf16_mm import ( # noqa: E501 + plan, +) -__extra_import_path__ = [".."] +__extra_import_path__ = ["../../../.."] from _validating import ( # noqa: E402,F401 — path above; fixture by name ULP, assert_within_ulp, diff --git a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py similarity index 93% rename from tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py index ca055aec966b..6750591cb2c1 100644 --- a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py @@ -1,20 +1,20 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the flashinfer_chunk_gated_delta_rule catalog entry (FlashInfer GDN prefill).""" +"""GPU test for the flashinfer_chunk_gated_delta_rule entry (FlashInfer GDN prefill).""" import zlib import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.flashinfer_chunk_gated_delta_rule import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.flashinfer_chunk_gated_delta_rule import ( # noqa: E501 QWEN38_GEOMETRY, ) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.flashinfer_chunk_gated_delta_rule import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.flashinfer_chunk_gated_delta_rule import ( # noqa: E501 flashinfer_chunk_gated_delta_rule as op, ) -__extra_import_path__ = [".."] +__extra_import_path__ = ["../../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "flashinfer_chunk_gated_delta_rule requires a CUDA device" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_in_proj_conv1d_pdl.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_in_proj_conv1d_pdl.py index a7c5c62612f4..6d5b2d1959b7 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_in_proj_conv1d_pdl.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_in_proj_conv1d_pdl.py @@ -16,7 +16,9 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.low_m_bf16_mm import low_m_bf16_mm +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.gemm.low_m_bf16_mm import ( # noqa: E501 + low_m_bf16_mm, +) 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, ) From 22c5e8776e830c4dc1082b17f1f24e77cc2b67c5 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Mon, 19 Oct 2026 19:46:09 -0700 Subject: [PATCH 17/28] [TRTLLM-17033][test] Qwen3.8-Flash-Next tests: assert the decode overlap, pin int64 positions, run the MoE scalars on CPU - test_breakable_prefill_matches_eager meant to prefill the 5000-token prompt while word problems decode, but submitted it 0.6 s after the first problem, by which time a step-by-step answer can be complete. The problems now go in first and the long prompt 0.2 s later, and the test asserts that a problem is still decoding when it arrives. - test_huge_int64_position_is_skipped_not_truncated builds its positions as int64 explicitly instead of through the first cell's dtype. - test_modeling_v2_qwen4_exp_moe_scalars derives the scalars from CPU tensors: it is marked cpu_only and listed in l0_cpu.yml, like the other CPU qwen4_exp files. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../accuracy/modeling_v2/test_qwen4_exp.py | 19 +++++++++---------- .../integration/test_lists/test-db/l0_cpu.yml | 5 +++-- .../test_modeling_v2_qsa_paged_index_store.py | 2 +- .../test_modeling_v2_qwen4_exp_moe_scalars.py | 2 ++ 4 files changed, 15 insertions(+), 13 deletions(-) diff --git a/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py b/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py index f24a4574a41f..cd4462bdb1b4 100644 --- a/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py +++ b/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py @@ -855,16 +855,15 @@ def run(engine): .token_ids for n in lengths ] - # The first requests are decoding when the long prompt arrives. - futures = [] - for index, prompt in enumerate(problems): - futures.append(llm.generate_async(prompt, answering)) - time.sleep(0.3) - if index == 1: - long_future = llm.generate_async( - {"prompt_token_ids": document[:5000]}, greedy - ) - time.sleep(0.3) + # The word problems are decoding when the long prompt arrives: their short + # prefills are done well inside the pause, and a step-by-step answer is + # still being written after it. + futures = [llm.generate_async(prompt, answering) for prompt in problems] + time.sleep(0.2) + assert not all(future.finished for future in futures), ( + "every word problem finished before the long prompt arrived" + ) + long_future = llm.generate_async({"prompt_token_ids": document[:5000]}, greedy) answers = [future.result().outputs[0].text for future in futures] long_tokens = long_future.result().outputs[0].token_ids return single, again, answers, long_tokens diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index ff2b78f2a921..6251b98757b7 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -70,8 +70,8 @@ l0_cpu: # compute capability, no device. The default targets' shard tests: small CPU # tensors and the CPU block_scale_interleave. The qwen4_exp entries are the target's # tests that run on the source text, stubbed weights and meta tensors: its multimodal - # splice and admission check on a stub tower, the KV pool contract, and its KV storage - # derivation and attention-call arguments. + # splice and admission check on a stub tower, the KV pool contract, its KV storage + # derivation and attention-call arguments, and its routed-expert scalars. # The rest of unittest/_torch/modeling_v2 needs a GPU. - unittest/_torch/modeling_v2/test_modeling_v2_claims.py - unittest/_torch/modeling_v2/test_modeling_v2_routing.py @@ -79,6 +79,7 @@ l0_cpu: - unittest/_torch/modeling_v2/test_modeling_v2_gpt_oss_default_shards.py - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_dense_backend_library_order - unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_multimodal_embeds.py + - unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_moe_scalars.py - unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_validate_request.py - unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_pool_contract.py - unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_kv_quant.py diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_store.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_store.py index 168be9db6969..151d027caab5 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_store.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_store.py @@ -409,7 +409,7 @@ def test_rejects_a_one_dimensional_table() -> None: def test_huge_int64_position_is_skipped_not_truncated() -> None: """`2^32 + 5` truncated to int32 would be position 5 of column 0; the kernel keeps int64 and finds it past the table.""" - args = _small(seed=70) + args = _small(seed=70, position_dtype=torch.int64) args["logical_positions"][0] = 2**32 + 5 before = (args["index_cache"].clone(), args["position_cache"].clone()) op(**args) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_moe_scalars.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_moe_scalars.py index a58994c365a7..a976733d0957 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_moe_scalars.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_moe_scalars.py @@ -19,6 +19,8 @@ modeling as M, ) +pytestmark = pytest.mark.cpu_only + EXPERTS = 8 From 37e2407335c201ae56065f35fe798bb3258ba0cc Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Mon, 19 Oct 2026 23:03:28 -0700 Subject: [PATCH 18/28] [TRTLLM-16460][chore] Qwen3.8-Flash-Next target: import the CuTe helpers from main's kernels tree #19275 moved tensorrt_llm/_torch/cute_dsl_kernels/blackwell/{utils,custom_pipeline}.py to tensorrt_llm/_torch/kernels/blackwell/ and the blockscaled grouped GEMM kernels to tensorrt_llm/_torch/moe/kernels/blackwell/. The target's CuTe ops import make_ptr and the pipeline helpers from the new modules, and the MoE ops' notes name the new paths. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../ops/attention/qsa_sparse_gqa_rowtile.py | 2 +- .../ops/attention/qwen4_exp_qk_norm_rope.py | 2 +- .../ops/hc/qwen4_exp_hc_decode_tc_first.py | 7 ++----- .../ops/hc/qwen4_exp_hc_inject_norm.py | 2 +- .../ops/hc/qwen4_exp_hc_up_gate_mean.py | 7 ++----- .../moe/cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py | 2 +- .../ops/moe/cute_dsl_nvfp4_grouped_gemm.py | 2 +- .../ops/ssm/causal_conv1d_update.py | 2 +- .../ops/ssm/gdn_decode_rows_gated_norm.py | 2 +- .../ops/ssm/gdn_prefill_conv_post.py | 2 +- 10 files changed, 12 insertions(+), 18 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py index 7251913ffade..bc4d8345f915 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py @@ -27,8 +27,8 @@ from cutlass.cutlass_dsl import T, dsl_user_op from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch.cute_dsl_kernels.blackwell.utils import make_ptr from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr from .qsa_paged_sparse_gqa import qsa_paged_sparse_gqa as _per_row_entry from .qsa_rowtile_attention_cute import BLOCK_M, BLOCK_N diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_qk_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_qk_norm_rope.py index 7e1d19411593..cf6f550b2e6c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_qk_norm_rope.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_qk_norm_rope.py @@ -27,7 +27,7 @@ from cutlass.cutlass_dsl import T, dsl_user_op from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch.cute_dsl_kernels.blackwell.utils import make_ptr +from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr @triton.jit diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_tc_first.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_tc_first.py index dfed495b80ca..22416db4bd12 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_tc_first.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_tc_first.py @@ -37,12 +37,9 @@ from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch.cute_dsl_kernels.blackwell.custom_pipeline import ( - PipelineTmaUmma, - PipelineUmmaAsync, -) -from tensorrt_llm._torch.cute_dsl_kernels.blackwell.utils import make_ptr from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch.kernels.blackwell.custom_pipeline import PipelineTmaUmma, PipelineUmmaAsync +from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr #: Largest row count this entry is certified for: one 128-row MMA tile. MAX_ROWS = 128 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_inject_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_inject_norm.py index 5b3b3cf75b28..6e61dbf47260 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_inject_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_inject_norm.py @@ -35,8 +35,8 @@ from cutlass.cutlass_dsl import T, dsl_user_op from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch.cute_dsl_kernels.blackwell.utils import make_ptr from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr #: Trigger dependents at kernel start (`griddepcontrol.launch_dependents`); a no-op #: unless the next kernel is launched with programmatic dependent launch. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean.py index 4e0e9a5eded2..8c681b9b8572 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean.py @@ -31,12 +31,9 @@ from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch.cute_dsl_kernels.blackwell.custom_pipeline import ( - PipelineTmaUmma, - PipelineUmmaAsync, -) -from tensorrt_llm._torch.cute_dsl_kernels.blackwell.utils import make_ptr from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch.kernels.blackwell.custom_pipeline import PipelineTmaUmma, PipelineUmmaAsync +from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr #: Programmatic dependent launch (sm_90+), off with TRTLLM_ENABLE_PDL=0. _PDL = get_env_enable_pdl() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py index a5bff838b059..8ab49bfa6aa6 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py @@ -214,7 +214,7 @@ class CuteDslNvfp4GatherGroupedGemmActFusion(OpWrapper): note = """ Path certified: the CuTe DSL kernel `BlockScaledContiguousGatherGroupedGemmKernel` - (`_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py`), + (`_torch/moe/kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py`), built by `Sm100BlockScaledContiguousGatherGroupedGemmActFusionRunner` (`_torch/custom_ops/cute_dsl_custom_ops.py`), which JIT-compiles it with `cute.compile` on the first call per (tile, tactic, top_k, activation) and diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_grouped_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_grouped_gemm.py index ed325ce10d52..c3398b0e579f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_grouped_gemm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_grouped_gemm.py @@ -207,7 +207,7 @@ class CuteDslNvfp4GroupedGemm(OpWrapper): note = """ Path certified: the CuTe DSL kernel `Sm100BlockScaledContiguousGroupedGemmKernel` - (`_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_grouped_gemm.py`), + (`_torch/moe/kernels/blackwell/blockscaled_contiguous_grouped_gemm.py`), built by `Sm100BlockScaledContiguousGroupedGemmRunner` (`_torch/custom_ops/cute_dsl_custom_ops.py`), which JIT-compiles it with `cute.compile` on the first call per (tile, tactic) and keeps it for the diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_update.py index f1d69faaaf0c..2e1846e5d478 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_update.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_update.py @@ -21,7 +21,7 @@ import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch.cute_dsl_kernels.blackwell.utils import make_ptr +from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr #: This checkpoint's GDN `[q|k|v]` width; the conv is width 4 over a 3-column pool. CKPT_DIM = 10240 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_rows_gated_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_rows_gated_norm.py index b81e768bb14a..9aed4c29e063 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_rows_gated_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_rows_gated_norm.py @@ -29,8 +29,8 @@ from cutlass.cutlass_dsl import T, dsl_user_op from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch.cute_dsl_kernels.blackwell.utils import make_ptr from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr from .gdn_decode_gated_norm_pdl import _recurrence_reference from .gdn_decode_gated_norm_pdl import gdn_decode_gated_norm_pdl as _one_launch diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_prefill_conv_post.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_prefill_conv_post.py index fe94663b51c0..29c567150675 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_prefill_conv_post.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_prefill_conv_post.py @@ -28,7 +28,7 @@ from cutlass.cutlass_dsl import T, dsl_user_op from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch.cute_dsl_kernels.blackwell.utils import make_ptr +from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr #: This checkpoint's linear-attention geometry at tp1. NUM_K_HEADS, HEAD_K_DIM, NUM_V_HEADS, HEAD_V_DIM = 16, 128, 48, 128 From ebae095a64d04636c9b48436b13d4aab39217b13 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Mon, 19 Oct 2026 23:41:17 -0700 Subject: [PATCH 19/28] [TRTLLM-16460][perf] Qwen3.8-Flash-Next target: one qsa_topk_expand_decode kernel for every row count The fused block top-k + expansion compiled one CuTe kernel per distinct row count, since its operands were bound with static layouts. The eager prefill path calls it with each step's context-row count, so a new count paid a JIT compile on the serving path. The entry now compiles through the TVM-FFI path the target's other CuTe ops use: the score table is passed as a tensor and its layout rebuilt inside with the static width and row stride the 16-byte key loads need, the other operands as pointers, and the row count as a runtime argument. The cache key drops the row count; is_valid also checks the per-row operands are contiguous, which the pointer form relies on. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../ops/attention/qsa_topk_expand_decode.py | 137 ++++++++++++------ ...test_modeling_v2_qsa_topk_expand_decode.py | 22 ++- 2 files changed, 112 insertions(+), 47 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_expand_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_expand_decode.py index f28662cc23fc..4d9bd57a18a3 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_expand_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_expand_decode.py @@ -19,18 +19,15 @@ from __future__ import annotations -import functools - -import cuda.bindings.driver as cuda_drv import cutlass import cutlass.cute as cute import torch from cutlass import const_expr -from cutlass.cute import experimental as cute_ext from cutlass.cute.runtime import from_dlpack from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr from .qsa_expand_block_indices import qsa_expand_block_indices as _expand @@ -60,9 +57,10 @@ def _threads(width: int) -> int: class _TopkExpandKernel: - def __init__(self, *, width, block_topk, compress, token_topk, pdl): - self.width, self.k, self.cr, self.token_topk, self.pdl = ( + def __init__(self, *, width, key_stride, block_topk, compress, token_topk, pdl): + self.width, self.key_stride, self.k, self.cr, self.token_topk, self.pdl = ( width, + key_stride, block_topk, compress, token_topk, @@ -74,8 +72,31 @@ def __init__(self, *, width, block_topk, compress, token_topk, pdl): self.flag_words = -(-width // 32) @cute.jit - def __call__(self, mKeys, mEnd, mPos, mLen, mTopk, mSel, stream: cuda_drv.CUstream): - rows = cute.size(mKeys, mode=[0]) + def __call__( + self, + keys_t: cute.Tensor, + end_ptr: cute.Pointer, + pos_ptr: cute.Pointer, + len_ptr: cute.Pointer, + topk_ptr: cute.Pointer, + sel_ptr: cute.Pointer, + rows: cutlass.Int32, + stream, + ): + # The keys arrive as a tensor, as the TVM-FFI entry takes its stream from a GPU + # tensor; their layout is rebuilt with the static width and row stride the 16-byte + # loads need. The row count is the call's. + mKeys = cute.make_tensor( + keys_t.iterator, cute.make_layout((rows, self.width), stride=(self.key_stride, 1)) + ) + per_row = cute.make_layout((rows,), stride=(1,)) + mEnd = cute.make_tensor(end_ptr, per_row) + mPos = cute.make_tensor(pos_ptr, per_row) + mLen = cute.make_tensor(len_ptr, per_row) + mTopk = cute.make_tensor(topk_ptr, cute.make_layout((rows, self.k), stride=(self.k, 1))) + mSel = cute.make_tensor( + sel_ptr, cute.make_layout((rows, self.final), stride=(self.final, 1)) + ) ld_keys = cute.make_copy_atom( cute.nvgpu.CopyG2ROp(), cutlass.Int32, num_bits_per_copy=32 * KPT ) @@ -367,26 +388,10 @@ def kernel(self, mKeys, mEnd, mPos, mLen, mTopk, mSel, ld_keys: cute.CopyAtom): mSel[row, col] = val -def _dl(t: torch.Tensor, align: int = 4): - return from_dlpack(t, assumed_align=align) - - -@functools.lru_cache(maxsize=None) -def _compiled(rows, width, key_stride, block_topk, compress, token_topk, pos_dtype, len_dtype, pdl): - kern = _TopkExpandKernel( - width=width, block_topk=block_topk, compress=compress, token_topk=token_topk, pdl=pdl - ) - dev = "cuda" - keys = torch.empty(rows, key_stride, dtype=torch.int32, device=dev)[:, :width] - end = torch.empty(rows, dtype=torch.int32, device=dev) - pos = torch.empty(rows, dtype=pos_dtype, device=dev) - ln = torch.empty(rows, dtype=len_dtype, device=dev) - topk = torch.empty(rows, block_topk, dtype=torch.int32, device=dev) - sel = torch.empty(rows, token_topk + compress - 1, dtype=torch.int32, device=dev) - stream = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream) - return cute_ext.compile( - kern, _dl(keys, 16), _dl(end), _dl(pos), _dl(ln), _dl(topk), _dl(sel), stream - ) +#: One compiled launch per score-table geometry, k, compress ratio, token budget, index +#: dtypes and PDL; the row count is a runtime argument. +_COMPILED: dict = {} +_INDEX_TYPES = {torch.int32: cutlass.Int32, torch.int64: cutlass.Int64} def launch( @@ -404,15 +409,16 @@ def launch( token selection `[rows, token_topk + compress_ratio - 1]` int32 (`out` when given).""" rows, width = scores.shape block_topk = topk_indices.shape[1] - dev = scores.device sel = out if sel is None: - sel = torch.empty((rows, token_topk + compress_ratio - 1), dtype=torch.int32, device=dev) + sel = torch.empty( + (rows, token_topk + compress_ratio - 1), dtype=torch.int32, device=scores.device + ) if rows == 0: return sel - pdl = pdl and get_env_enable_pdl() - fn = _compiled( - rows, + pdl = bool(pdl and get_env_enable_pdl()) + keys = scores.view(torch.int32) + key = ( width, scores.stride(0), block_topk, @@ -420,18 +426,48 @@ def launch( token_topk, query_positions.dtype, sequence_lengths.dtype, - bool(pdl), + pdl, ) - stream = cuda_drv.CUstream(torch.cuda.current_stream(dev).cuda_stream) - fn( - _dl(scores.view(torch.int32), 16), - _dl(row_ends), - _dl(query_positions), - _dl(sequence_lengths), - _dl(topk_indices), - _dl(sel), - stream, + ptrs = ( + row_ends.data_ptr(), + query_positions.data_ptr(), + sequence_lengths.data_ptr(), + topk_indices.data_ptr(), + sel.data_ptr(), ) + fn = _COMPILED.get(key) + if fn is None: + gm = cute.AddressSpace.gmem + fn = _COMPILED[key] = cute.compile( + _TopkExpandKernel( + width=width, + key_stride=scores.stride(0), + block_topk=block_topk, + compress=compress_ratio, + token_topk=token_topk, + pdl=pdl, + ), + from_dlpack(keys, assumed_align=16).mark_layout_dynamic(leading_dim=1), + make_ptr(cutlass.Int32, ptrs[0], gm, assumed_align=4), + make_ptr( + _INDEX_TYPES[query_positions.dtype], + ptrs[1], + gm, + assumed_align=query_positions.element_size(), + ), + make_ptr( + _INDEX_TYPES[sequence_lengths.dtype], + ptrs[2], + gm, + assumed_align=sequence_lengths.element_size(), + ), + make_ptr(cutlass.Int32, ptrs[3], gm, assumed_align=4), + make_ptr(cutlass.Int32, ptrs[4], gm, assumed_align=4), + cutlass.Int32(rows), + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + fn(keys, *ptrs, rows) return sel @@ -495,8 +531,13 @@ class QsaTopkExpandDecode(OpWrapper): Metadata consumed: `qsa_visible_blocks` (`row_ends`), the rows' logical positions and sequence lengths. Preconditions: zero row starts; `scores` fp32 `[rows, W]`, unit column stride, 16-byte aligned rows, `W % 4 == 0` and `W <= 4096` (four keys a thread, - at most 1024 threads); `topk_indices` int32 `[rows, K]` contiguous. Ties: the higher + at most 1024 threads); `topk_indices` int32 `[rows, K]` contiguous; `row_ends` int32, + the positions and lengths int32 or int64, each a contiguous `[rows]`. Ties: the higher block at the threshold, as `qsa_canonical_selection` rewrites the CUDA top-k's pick. + + Compiled once per score-table width and row stride (the block table's, fixed for an + engine), k, compress ratio, token budget, index dtypes and PDL: the row count is a + runtime argument, so an eager step's context rows, any count, reuse the kernel. """ def __call__( @@ -574,8 +615,12 @@ def is_valid( assert topk_indices.dtype == torch.int32 and topk_indices.is_contiguous(), ( "int32 topk_indices" ) - assert row_ends.dtype == torch.int32 and row_ends.shape == (scores.shape[0],), ( - "int32 row_ends [rows]" + rows = scores.shape[0] + assert row_ends.dtype == torch.int32 and row_ends.shape == (rows,), "int32 row_ends [rows]" + for t in (query_positions, sequence_lengths): + assert t.dtype in _INDEX_TYPES and t.shape == (rows,), "int32/int64 [rows]" + assert all(t.is_contiguous() for t in (row_ends, query_positions, sequence_lengths)), ( + "contiguous row_ends, positions and lengths" ) assert token_topk == topk_indices.shape[1] * compress_ratio, ( "token_topk = k * compress_ratio" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_expand_decode.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_expand_decode.py index d13a4572f629..966013ce94d0 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_expand_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_expand_decode.py @@ -7,7 +7,7 @@ Columns past each row's visible count carry garbage (huge values, NaN) that must never be read. Beyond the cells: bitwise the composition it replaces (`qsa_topk_select`, then `qsa_canonical_selection`, then `qsa_expand_block_indices`) on tie-free and tie-heavy rows; -a stale-read check under graph replay. +a stale-read check under graph replay; one compiled kernel for every row count. """ import zlib @@ -21,6 +21,9 @@ 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, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_topk_expand_decode import ( # noqa: E501 + _COMPILED, +) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_topk_expand_decode import ( # noqa: E501 qsa_topk_expand_decode as op, ) @@ -162,3 +165,20 @@ def step(): got_topk = args[4].clone() sel = op(src.clone(), *args[1:]) assert torch.equal(graph_sel, sel) and torch.equal(got_topk, args[4]), f"replay {trial}" + + +def test_one_kernel_serves_every_row_count() -> None: + """The row count is a runtime argument: an eager step's context rows come in any count, + and every count here runs the one kernel compiled for the table, each matching the + reference.""" + compiled = None + for rows in (1, 2, 7, 33, 300): + g = torch.Generator(device=DEV).manual_seed(rows) + pos = torch.randint(0, CR * W, (rows,), generator=g, device=DEV, dtype=torch.int64) + args = _build(dict(pos=tuple(pos.tolist())), seed=rows) + want_topk, want_sel = op.reference(*args) + sel = op(*args) + assert torch.equal(args[4], want_topk), f"block list, {rows} rows" + assert torch.equal(sel, want_sel.to(torch.int32)), f"token selection, {rows} rows" + compiled = len(_COMPILED) if compiled is None else compiled + assert len(_COMPILED) == compiled, f"{rows} rows compiled another kernel" From 65f7c68376161ee3507116ad298e16094a1a1eca Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Wed, 21 Oct 2026 00:23:22 -0700 Subject: [PATCH 20/28] [TRTLLM-17033][fix] thop attention: FP8 context FMHA for a non-MLA NVFP4 KV cache only on SM100f Only the SM100-family FP8 context FMHA kernels read an NVFP4 KV cache, so other architectures keep failing AttentionOp's FP4 KV check instead of passing the FP8 FMHA gate. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- cpp/tensorrt_llm/thop/attentionOp.cpp | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/cpp/tensorrt_llm/thop/attentionOp.cpp b/cpp/tensorrt_llm/thop/attentionOp.cpp index f588dd465a45..b1089b34134f 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.cpp +++ b/cpp/tensorrt_llm/thop/attentionOp.cpp @@ -1355,10 +1355,12 @@ void attention(torch::Tensor q, std::optional k, std::optionalmRotaryEmbeddingLongMscale = rope_long_m_scale; op->mRotaryEmbeddingMaxPositions = rope_max_positions; op->mRotaryEmbeddingOriginalMaxPositions = rope_original_max_positions; - // Only the FP8 context FMHA kernels read an NVFP4 KV cache, whatever the output type. MLA keeps its own NVFP4 KV + // Only the FP8 context FMHA kernels read an NVFP4 KV cache, whatever the output type; they do so on SM100-family + // GPUs only, so elsewhere an NVFP4 KV cache still fails the FP4 KV check in AttentionOp. MLA keeps its own NVFP4 KV // handling (mUseNvfp4MlaKvCache, mFP8ContextMLA below) and must not pick this up. op->mFP8ContextFMHA = is_fp8_out || is_fp4_out || (op->mKVCacheQuantMode.hasFp8KvCache() && use_paged_context_fmha) - || (!is_mla_enable && op->mKVCacheQuantMode.hasFp4KvCache()) || use_sage_attn; + || (!is_mla_enable && op->mKVCacheQuantMode.hasFp4KvCache() && tensorrt_llm::common::isSM100Family()) + || use_sage_attn; // SageAttention block sizes and quantization mode. op->mSageAttnNumEltsPerBlkQ = static_cast(sage_attn_num_elts_per_blk_q); op->mSageAttnNumEltsPerBlkK = static_cast(sage_attn_num_elts_per_blk_k); From bf4dbebf1e22a720c2f21909f33756f2d94ff3cc Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Wed, 21 Oct 2026 00:23:26 -0700 Subject: [PATCH 21/28] [TRTLLM-16460][perf] Qwen3.8-Flash-Next target: keep the second expert copy only where its launches can happen Under trtllm-gen the trunk kept a CuteDSL-layout copy of the routed experts on every layer (about 63 GiB) for context launches of at least 4096 rows. It is now allocated and dispatched only when max_num_tokens can reach that size, and TRTLLM_MODELING_V2_MOE_PREFILL_COPY=0 drops it for memory-constrained deployments; every launch then runs on trtllm-gen. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../modeling.py | 38 ++++++++++++--- .../weights.py | 11 +++-- ...test_modeling_v2_qwen4_exp_weight_table.py | 48 +++++++++++++++---- 3 files changed, 77 insertions(+), 20 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py index 49c25e221fa4..1da24685c9b5 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py @@ -45,7 +45,9 @@ chosen by ``moe_config.backend``; any other request is refused. The MTP draft layer's FP8 block-scale experts run on trtllm-gen under either. Under ``TRTLLM`` the prefill-sized launches (``_CUTE_DSL_MOE_MIN_ROWS``) run the CuteDSL pipeline over a - second copy of the experts in its layout. A graph-captured decode step routes its block + second copy of the experts in its layout, about 1.3 GiB a layer, kept only when a step + can reach that size and ``TRTLLM_MODELING_V2_MOE_PREFILL_COPY`` does not turn it off + (``_moe_prefill_copy``). A graph-captured decode step routes its block projections through ``ops/gemm/low_m_bf16_mm`` (FlashInfer low-M kernels at M <= 32 on the shapes it lists, cuBLAS otherwise); see ``DecodeTarget.forward``. @@ -503,6 +505,24 @@ def _ple_host_offload_enabled() -> bool: return os.environ.get(_PLE_HOST_OFFLOAD_ENV, "1").strip().lower() not in _PLE_HOST_OFFLOAD_OFF +# Whether trtllm-gen keeps the second, CuteDSL-layout copy of the routed experts for the +# launches of at least `_CUTE_DSL_MOE_MIN_ROWS` rows (module docstring). No engine arg owns +# it, hence an env var read at construction; "off" (the spellings above) drops the copy and +# every launch runs on trtllm-gen. +_MOE_PREFILL_COPY_ENV = "TRTLLM_MODELING_V2_MOE_PREFILL_COPY" + + +def _moe_prefill_copy(trunk_moe_backend: str, max_num_tokens: int) -> bool: + """Whether the trunk keeps the second expert copy: only under trtllm-gen, only when a + step's rows (at most `max_num_tokens`) can reach the CuteDSL dispatch threshold, and not + when the env var turns it off.""" + return ( + trunk_moe_backend == "TRTLLM" + and max_num_tokens >= _CUTE_DSL_MOE_MIN_ROWS + and os.environ.get(_MOE_PREFILL_COPY_ENV, "1").strip().lower() not in _PLE_HOST_OFFLOAD_OFF + ) + + # QSA's fixed constants (backend `constants.py`): one index KV head, three position # axes, index head exactly twice the rotary width (128 == 2 * 64). _QSA_INDEX_KV_HEADS = 1 @@ -1256,6 +1276,9 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig f"this target; the supported requests are {sorted(_TRUNK_MOE_BACKENDS)}" ) self.trunk_moe_backend = _TRUNK_MOE_BACKENDS[requested] + self.moe_prefill_copy = _moe_prefill_copy( + self.trunk_moe_backend, model_config.max_num_tokens + ) # ---- PLE ----------------------------------------------------------- ple_ids = list(getattr(cfg, "ple_layer_ids", None) or []) @@ -1556,11 +1579,11 @@ def post_load_weights(self) -> None: moe_scalars = [ _moe_layer_scalars(w, i, self.trunk_moe_backend) for i in range(self.num_layers) ] - # Under trtllm-gen, the CuteDSL scalars of the second expert copy that + # With the second expert copy (`moe_prefill_copy`), its CuteDSL scalars, which # prefill-sized launches run on (`_CUTE_DSL_MOE_MIN_ROWS`). moe_scalars_cd = ( [_moe_layer_scalars(w, i, "CUTEDSL") for i in range(self.num_layers)] - if self.trunk_moe_backend == "TRTLLM" + if self.moe_prefill_copy else None ) self._moe_scalars = moe_scalars @@ -2081,7 +2104,8 @@ class _LayerWeights(NamedTuple): alpha2: torch.Tensor | None global_sf: torch.Tensor | None # Under trtllm-gen, the second expert copy in the CuteDSL layout and its scalars, - # which prefill-sized launches run on (`_CUTE_DSL_MOE_MIN_ROWS`); None on CuteDSL. + # which prefill-sized launches run on (`_CUTE_DSL_MOE_MIN_ROWS`); None on CuteDSL and + # without the copy (`moe_prefill_copy`). fc1_w_cd: torch.Tensor | None fc1_s_cd: torch.Tensor | None fc2_w_cd: torch.Tensor | None @@ -2756,8 +2780,8 @@ def _moe_block(self, x, xq, xsf, L: _LayerWeights, routed_out) -> tuple: runner the router logits (its Renormalize routing is bitwise the `renorm_moe_routing_op` + pre-routed composition), and a launch of at least `_CUTE_DSL_MOE_MIN_ROWS` rows runs the CuteDSL MoE over the second - expert copy instead; on CuteDSL every launch is `_cute_dsl_routed`, which - writes `routed_out`.""" + expert copy instead when the core keeps it (`moe_prefill_copy`); on CuteDSL + every launch is `_cute_dsl_routed`, which writes `routed_out`.""" core = self.core main = torch.cuda.current_stream() side = core._side_stream @@ -2777,7 +2801,7 @@ def _moe_block(self, x, xq, xsf, L: _LayerWeights, routed_out) -> tuple: logits = cublas_mm(x, L.router) if core.trunk_moe_backend == "CUTEDSL": routed = _cute_dsl_routed(core, L, logits, xq, xsf, routed_out) - elif x.shape[0] >= _CUTE_DSL_MOE_MIN_ROWS: + elif core.moe_prefill_copy and x.shape[0] >= _CUTE_DSL_MOE_MIN_ROWS: routed = fp4_moe_renorm_cute_dsl( logits, xq, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/weights.py index 83662a561062..1df6b9943f24 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/weights.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/weights.py @@ -22,7 +22,8 @@ interleaves FC1 row by row and 32-row block shuffles weight and scale bytes (FC2: shuffle only); CuteDSL interleaves FC1 in 64-row up/gate groups and leaves FC2 in checkpoint row order. Under trtllm-gen a second copy in the CuteDSL layout (`fc*_cd`) - serves the prefill-sized launches (`modeling._CUTE_DSL_MOE_MIN_ROWS`); + serves the prefill-sized launches (`modeling._CUTE_DSL_MOE_MIN_ROWS`) when the core keeps + it (`core.moe_prefill_copy`); * n-gram table: 128 row chunks concatenated into one `[rows, 160]` buffer (the hash emits global row indices). @@ -425,9 +426,10 @@ def _in_proj_rows(d, part: int) -> tuple[slice]: def _cute_dsl_copy_layers(d): - """Every layer when the trunk's routed experts run on trtllm-gen (the second, CuteDSL - copy that prefill-sized launches read); none under CuteDSL, whose one copy is it.""" - return range(d.num_layers) if d.trunk_moe_backend == "TRTLLM" else () + """Every layer when the core keeps the second, CuteDSL-layout copy that prefill-sized + launches read under trtllm-gen (`core.moe_prefill_copy`); none otherwise, CuteDSL among + them, whose one copy is it.""" + return range(d.num_layers) if d.moe_prefill_copy else () def _vision_blocks(d): @@ -1158,6 +1160,7 @@ def dims(self, core) -> SimpleNamespace: v_merge_dim=v_hidden * vcfg.spatial_merge_size * vcfg.spatial_merge_size, mtp_enabled=getattr(core.model_config, "spec_config", None) is not None, trunk_moe_backend=core.trunk_moe_backend, + moe_prefill_copy=core.moe_prefill_copy, dtype=core.model_config.torch_dtype, ) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_weight_table.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_weight_table.py index 60fd360c0352..8eea4907b6c4 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_weight_table.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_weight_table.py @@ -20,6 +20,9 @@ __extra_import_path__ = [".."] # noqa: F841 -- repo's file-scoped import hook +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 + modeling as M, +) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 weights as W, ) @@ -65,8 +68,9 @@ ) -def _core(mtp: bool, backend: str = "TRTLLM"): - """The configuration surface the table reads.""" +def _core(mtp: bool, backend: str = "TRTLLM", prefill_copy: bool | None = None): + """The configuration surface the table reads; the second expert copy defaults to kept + under trtllm-gen.""" return SimpleNamespace( model_config=SimpleNamespace( pretrained_config=SimpleNamespace(**_TEXT), @@ -76,6 +80,7 @@ def _core(mtp: bool, backend: str = "TRTLLM"): vision_config=SimpleNamespace(**_VISION), mtp_enabled=mtp, trunk_moe_backend=backend, + moe_prefill_copy=backend == "TRTLLM" if prefill_copy is None else prefill_copy, ) @@ -361,12 +366,37 @@ def test_cutedsl_fc2_keeps_the_checkpoint_row_order_and_trtllm_gen_shuffles_it() assert int(shuffled[32 + 8 * v + u, 0]) == 32 + 4 * u + v -@pytest.mark.parametrize("backend", ["TRTLLM", "CUTEDSL"]) -def test_cute_dsl_expert_copy_rows_only_under_trtllm_gen(backend): - """Under trtllm-gen every layer also carries its experts in the CuteDSL layout - (`fc*_cd`), which prefill-sized launches run on; under CuteDSL the one copy is that - layout and no `_cd` row is declared.""" - manifest = W.MODEL_WEIGHTS.manifest(_core(False, backend)) +@pytest.mark.parametrize( + "backend, prefill_copy", [("TRTLLM", True), ("TRTLLM", False), ("CUTEDSL", False)] +) +def test_cute_dsl_expert_copy_rows_only_when_kept(backend, prefill_copy): + """With the second expert copy (trtllm-gen, a step able to reach the CuteDSL dispatch + threshold) every layer also carries its experts in the CuteDSL layout (`fc*_cd`); + without it, and under CuteDSL whose one copy is that layout, no `_cd` row is declared.""" + manifest = W.MODEL_WEIGHTS.manifest(_core(False, backend, prefill_copy)) rows = [k for k in manifest if k.endswith(("_fc1_w_cd", "_fc1_s_cd", "_fc2_w_cd", "_fc2_s_cd"))] layers = _TEXT["num_hidden_layers"] - assert len(rows) == (4 * layers if backend == "TRTLLM" else 0), len(rows) + assert len(rows) == (4 * layers if prefill_copy else 0), len(rows) + + +@pytest.mark.parametrize( + "backend, max_num_tokens, env, kept", + [ + ("TRTLLM", 16384, None, True), + ("TRTLLM", 4096, None, True), + ("TRTLLM", 4095, None, False), + ("TRTLLM", 16384, "0", False), + ("TRTLLM", 16384, "1", True), + ("CUTEDSL", 16384, None, False), + ], +) +def test_second_expert_copy_kept_only_where_its_launches_can_happen( + monkeypatch, backend, max_num_tokens, env, kept +): + """The copy costs about 1.3 GiB a layer: it is kept only under trtllm-gen, only when a + step's rows can reach the CuteDSL dispatch threshold, and not when the env var is off.""" + if env is None: + monkeypatch.delenv(M._MOE_PREFILL_COPY_ENV, raising=False) + else: + monkeypatch.setenv(M._MOE_PREFILL_COPY_ENV, env) + assert M._moe_prefill_copy(backend, max_num_tokens) is kept From b3e6c7531849e0f4a50076bd4c6e453cd4695697 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Wed, 21 Oct 2026 00:23:28 -0700 Subject: [PATCH 22/28] [TRTLLM-17033][fix] QSA cache manager: side-plane transfer roles independent of the KV head split The index and position planes were packed into QSA_SIDE_PLANES when a page coincided in size with a K/V or block-scale page on this rank, which depends on the KV heads the rank holds: with FP8 KV, two KV heads, head_dim 256 and 32 tokens per block, TP1 registered INDEX_KEY plus QSA_INDEX_POSITION while TP2 (one head a rank) registered the packed role, and NVFP4 flipped which side packed against ADP. Disaggregated peers then failed validate_peer_compatible. The candidate page sizes are now every size a layer's K/V or block-scale page can take at any split of its KV heads over ranks (one head a rank up to all of them), so the packing decision and the packed size are the same on TP, TP>1 and attention-DP layouts. With this target's geometry FP8 and NVFP4 pack on every layout and bf16 on none. The packed buffer grows in 16-byte steps. Adds a TP1/TP2/ADP2 regression test on real managers: equal logical roles, mapper kinds and side-view pages, with physical placement free to differ. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../backends/sparse/qsa/cache_manager.py | 49 +++++++---- .../qsa/test_qsa_cache_manager_layout.py | 87 +++++++++++++++---- .../attention/sparse/qsa/test_qsa_sparse.py | 4 +- 3 files changed, 106 insertions(+), 34 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py index 442cb8d67758..0d5273ebef92 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py @@ -36,17 +36,17 @@ # 16-element block. _NVFP4_VALUES_PER_BYTE = 2 _NVFP4_BLOCK_SIZE = 16 -# Granularity by which the packed side-plane page grows until its size is -# distinct from every main-cache page; keeps each plane 16-byte aligned. -_SIDE_PLANES_SIZE_STEP = 128 +# Granularity by which the packed side-plane page grows until its size cannot +# be a main-cache page; keeps every slot's planes 16-byte aligned. +_SIDE_PLANES_SIZE_STEP = 16 # Per-token RoPE/mRoPE coordinates used when raw index keys are compressed. # They are request state shared by all local QSA layers. QSA_INDEX_POSITION = DataRole("qsa_index_position") # Every local layer's index-K page followed by the position page, as one # buffer per lifecycle slot. Registered instead of the per-layer INDEX_KEY and -# the position role when either page would be the size of a K/V or -# block-scale page (see `_extra_buffers_per_layer`). +# the position role when either page could be the size of a K/V or block-scale +# page (see `_extra_buffers_per_layer`). QSA_SIDE_PLANES = DataRole("qsa_side_planes") @@ -65,9 +65,12 @@ class QSAMambaHybridCacheManagerV2(MambaHybridCacheManagerV2): attention op addresses a layer's K/V pages and its NVFP4 block-scale pages with one block-offset table, so a side page that lands in the K/V pool (NVFP4 at 32 tokens per block: index-K 8192 B == K == V) changes the K/V - slot stride but not the block-scale one. The side planes are then - registered as one packed per-slot buffer whose size matches no main-cache - page, and this manager serves their per-layer views itself. + slot stride but not the block-scale one, and any side page in the K/V pool + breaks the K/V pool mapping. When a side page could be the size of a main + page, the side planes are therefore registered as one packed per-slot buffer + whose size can be no main-cache page, and this manager serves their + per-layer views itself. Both the choice and the packed size are the same on + every TP and attention-DP layout, so disaggregated peers agree on the roles. """ def __init__( @@ -137,11 +140,14 @@ def _extra_buffers_per_layer( ) -> dict[int, list[BufferConfig]]: """Register per-layer index K and lifecycle-aligned position pages. - Both are registered as their own roles unless either page has the - size of a K/V or block-scale page, which V2 would coalesce into that - pool. They are then packed, in local sparse-layer order with the - position page last, into one `QSA_SIDE_PLANES` buffer on the first - local sparse layer, grown until its size matches no main-cache page. + Both are registered as their own roles unless either page could have + the size of a K/V or block-scale page at some split of the KV heads over + ranks (`_main_cache_page_sizes`), which V2 would coalesce into that pool. + They are then packed, in local sparse-layer order with the position page + last, into one `QSA_SIDE_PLANES` buffer on the first local sparse layer, + grown until its size is none of those page sizes. Neither the choice nor + the size depends on the KV heads this rank holds, so peers of different + TP or attention-DP layouts register the same transfer roles. """ index_size = ( self.qsa_index_kv_heads * self.qsa_index_dim * _INDEX_K_ELEMENT_BYTES * tokens_per_block @@ -171,13 +177,22 @@ def _extra_buffers_per_layer( return result def _main_cache_page_sizes(self, tokens_per_block: int) -> set[int]: - """Page sizes of every local K/V and block-scale buffer, in bytes.""" + """Every size, in bytes, a local layer's K/V or block-scale page can have + at any split of that layer's KV heads over ranks: one head per rank up to + all of them. The set is the same on every TP and attention-DP layout.""" + global_layer = {local: layer for layer, local in self.layer_offsets.items()} sizes = set() for local_idx in range(self.num_local_layers): + heads = self.num_kv_heads_per_layer[local_idx] + total = ( + self.num_kv_heads + if isinstance(self.num_kv_heads, int) + else self.num_kv_heads[global_layer[local_idx]] + ) for role in self._get_buffer_roles_for_layer(local_idx): - size = self.get_layer_bytes_per_token(local_idx, role) * tokens_per_block - if size: - sizes.add(size) + page = self.get_layer_bytes_per_token(local_idx, role) * tokens_per_block + if page: + sizes.update(page // heads * h for h in range(1, total + 1)) return sizes def _side_planes(self) -> torch.Tensor: diff --git a/tests/unittest/_torch/attention/sparse/qsa/test_qsa_cache_manager_layout.py b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_cache_manager_layout.py index 5c811ea87d5b..18abace17337 100644 --- a/tests/unittest/_torch/attention/sparse/qsa/test_qsa_cache_manager_layout.py +++ b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_cache_manager_layout.py @@ -42,8 +42,11 @@ _POSITION_PAGE = 3 * 4 * _TPB -def _manager(dtype: DataType, *, draft: bool) -> QSAMambaHybridCacheManagerV2: - """The trunk or the separate MTP draft manager, as the engine builds them.""" +def _manager( + dtype: DataType, *, draft: bool, tp_size: int = 1, attention_dp: bool = False +) -> QSAMambaHybridCacheManagerV2: + """The trunk or the separate MTP draft manager, as the engine builds them (rank 0 of + `tp_size`, with attention data parallelism when `attention_dp`).""" if draft: attention_mask = [False] * _TRUNK_LAYERS + [True] mamba_mask = [False] * (_TRUNK_LAYERS + 1) @@ -68,7 +71,13 @@ def _manager(dtype: DataType, *, draft: bool) -> QSAMambaHybridCacheManagerV2: tokens_per_block=_TPB, max_seq_len=1024, max_batch_size=4, - mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + mapping=Mapping( + world_size=tp_size, + rank=0, + tp_size=tp_size, + pp_size=1, + enable_attention_dp=attention_dp, + ), dtype=dtype, layer_mask=attention_mask, is_draft=draft, @@ -145,18 +154,21 @@ def test_views_are_where_the_attention_op_addresses(dtype: DataType, draft: bool @pytest.mark.parametrize("draft", [False, True], ids=["trunk", "draft"]) -def test_nvfp4_scale_pages_follow_the_data_pages(draft: bool) -> None: - """Under NVFP4 the K/V pool and the block-scale pool hold the same number - of pages per slot, and the index and position planes are packed into a - buffer of their own instead of landing in the K/V pool.""" - cm = _manager(DataType.NVFP4, draft=draft) +@pytest.mark.parametrize("dtype", [DataType.FP8, DataType.NVFP4]) +def test_side_planes_pack_when_a_page_could_collide(dtype: DataType, draft: bool) -> None: + """An index page could be a K/V page here -- FP8 at one KV head a rank (1 x 256 x 32 + tokens), NVFP4 at both -- so the index and position planes are packed into a buffer + of their own instead of landing in the K/V pool, and under NVFP4 the K/V and + block-scale pools hold the same number of pages per slot.""" + cm = _manager(dtype, draft=draft) try: layers = _layers(draft) assert cm.qsa_side_planes_size == len(layers) * _INDEX_PAGE + _POSITION_PAGE for ordinal, layer in enumerate(layers): local = cm.layer_offsets[layer] data_scale = cm.impl.get_page_index_scale(local, Role.KEY) - assert data_scale == cm.impl.get_page_index_scale(local, Role.KEY_BLOCK_SCALE) + if dtype == DataType.NVFP4: + assert data_scale == cm.impl.get_page_index_scale(local, Role.KEY_BLOCK_SCALE) assert data_scale == 2 * len(layers) assert cm.kv_cache_pool_mapping[local, 1] == ordinal planes_owner = cm.layer_offsets[layers[0]] @@ -166,12 +178,11 @@ def test_nvfp4_scale_pages_follow_the_data_pages(draft: bool) -> None: @pytest.mark.parametrize("draft", [False, True], ids=["trunk", "draft"]) -@pytest.mark.parametrize("dtype", [DataType.BF16, DataType.FP8]) -def test_bf16_and_fp8_keep_per_layer_side_roles(dtype: DataType, draft: bool) -> None: - """bf16 and fp8 pages never share a size with an index page, so the side - planes stay the per-layer INDEX_KEY role plus one position role, served by - the same views the generic accessor builds.""" - cm = _manager(dtype, draft=draft) +def test_bf16_keeps_per_layer_side_roles(draft: bool) -> None: + """A bf16 K/V page (at least 1 x 256 x 2 bytes x 32 tokens) can never be an index or + position page, so the side planes stay the per-layer INDEX_KEY role plus one position + role, served by the same views the generic accessor builds.""" + cm = _manager(DataType.BF16, draft=draft) try: layers = _layers(draft) assert cm.qsa_side_planes_size is None @@ -222,3 +233,49 @@ def test_side_and_main_views_do_not_alias(dtype: DataType, draft: bool) -> None: assert bool((raw == i + 1).all()), f"view {i} was overwritten by another view" finally: cm.shutdown() + + +def _logical_page(view: torch.Tensor) -> tuple: + """What a peer copies of one side view per slot: the page's shape, dtype and bytes, and + whether they are contiguous. Where the page sits in its pool (the slot stride, other + buffers coalesced beside it) is physical and may differ between layouts.""" + page = view[0] + return tuple(page.shape), page.dtype, page.numel() * page.element_size(), page.is_contiguous() + + +#: (tp_size, attention_dp): the layouts a disaggregated pair can mix over one global +#: geometry (2 KV heads): TP1 and ADP2 hold both heads per rank, TP2 one. +_LAYOUTS = {"tp1": (1, False), "tp2": (2, False), "adp2": (2, True)} + + +@pytest.mark.parametrize("dtype", [DataType.BF16, DataType.FP8, DataType.NVFP4]) +def test_transfer_views_do_not_depend_on_the_kv_head_split(dtype: DataType) -> None: + """Disaggregated peers compare the replicated side roles each layer registers and copy + their pages byte for byte, so roles, mapper kinds and every side view's page (shape, + dtype, bytes) must not depend on how many KV heads a rank holds; where pages sit in + their pools may. At TP2 an FP8 K/V page (1 head x 256 x 32 tokens) is the size of an + index page; under NVFP4 the same holds at TP1 and ADP2 (2 heads x 128 bytes x 32 + tokens). Both pack on every layout; bf16, which can never collide, on none.""" + seen = {} + for name, (tp_size, attention_dp) in _LAYOUTS.items(): + cm = _manager(dtype, draft=False, tp_size=tp_size, attention_dp=attention_dp) + try: + buffers = cm._extra_buffers_per_layer(tokens_per_block=_TPB) + roles = { + layer: sorted( + (str(b.role), b.size) for b in buffers.get(cm.layer_offsets[layer], []) + ) + for layer in _QSA_LAYERS + } + views = {layer: cm.get_index_k_buffer(layer) for layer in _QSA_LAYERS} + views["position"] = cm.get_qsa_position_buffer() + logical = {key: _logical_page(view) for key, view in views.items()} + seen[name] = (roles, logical, cm.get_disagg_role_mapper_kinds()) + finally: + cm.shutdown() + assert seen["tp1"] == seen["tp2"] == seen["adp2"] + roles = {role for layer in seen["tp1"][0].values() for role, _ in layer} + if dtype == DataType.BF16: + assert roles == {str(Role.INDEX_KEY), str(QSA_INDEX_POSITION)} + else: + assert roles == {str(QSA_SIDE_PLANES)} diff --git a/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py index ddc36f14be14..59308c3778db 100644 --- a/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py +++ b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py @@ -1520,8 +1520,8 @@ def test_qsa_side_buffers_pack_when_a_page_size_matches_the_main_cache() -> None assert list(buffers) == [2] (packed_buffer,) = buffers[2] assert packed_buffer.role == QSA_SIDE_PLANES - assert packed_buffer.size == packed + 128 - assert manager.qsa_side_planes_size == packed + 128 + assert packed_buffer.size == packed + 16 + assert manager.qsa_side_planes_size == packed + 16 def test_qsa_position_buffer_keeps_shared_index_mode_outside_dynamo() -> None: From a8ce4e549f78a0e4cc09678febf4c8bda912a75a Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Wed, 21 Oct 2026 02:10:51 -0700 Subject: [PATCH 23/28] [TRTLLM-17033][test] Qwen3.8-Flash-Next accuracy gates: 26 cases down to 11 The four pre-merge gates are unchanged. Post-merge keeps seven: three drafts over CuteDSL experts and a bf16 recurrent state in one engine (merging the two separate gates), the FP8 and NVFP4 KV pass keys, seven drafts over 3k-12k-token pass keys now on a bf16 state (merging the fp32 and bf16 versions), seven drafts all rejected, video, and the breakable prefill graph against eager. The thirteen cases no list ran are removed, with the acceptance-length keys only they read. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../accuracy/modeling_v2/test_qwen4_exp.py | 356 +++--------------- .../references/acceptance_length.yaml | 97 +---- .../test_lists/test-db/l0_b300.yml | 21 +- 3 files changed, 71 insertions(+), 403 deletions(-) diff --git a/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py b/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py index cd4462bdb1b4..45aa42834bea 100644 --- a/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py +++ b/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py @@ -412,33 +412,26 @@ class TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1(LlmapiAccuracyTestHarness): moe_config=MOE, ) - # The breakable prefill CUDA graph is opt-in. The buckets cover the gsm8k prompts; with the - # 2048-token budget every chunk of a longer prompt lands on one as well. + # The breakable prefill CUDA graph is opt-in; these buckets cover the prompts of + # `test_breakable_prefill_matches_eager`. BREAKABLE_PREFILL = dict( prefill_cuda_graph_backend="breakable", prefill_capture_num_tokens=[128, 256, 512, 1024, 2048, 3072], ) ENGINE_BREAKABLE = dict(ENGINE, **BREAKABLE_PREFILL) - ENGINE_BREAKABLE_CHUNKED = dict(ENGINE_BREAKABLE, max_num_tokens=2048) # One-model MTP-Eagle with the checkpoint's single MTP module replayed # `max_draft_len` times. MTP1 = MTPDecodingConfig(max_draft_len=1) MTP3 = MTPDecodingConfig(max_draft_len=3) - MTP5 = MTPDecodingConfig(max_draft_len=5) MTP7 = MTPDecodingConfig(max_draft_len=7) # Acceptance-length anchors in references/acceptance_length.yaml, keyed by # target and variant. ACCEPTANCE_KEY = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp1" - ACCEPTANCE_KEY_FP8_KV = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp1_fp8_kv" - ACCEPTANCE_KEY_NVFP4_KV = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp1_nvfp4_kv" - ACCEPTANCE_KEY_MTP3 = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3" - ACCEPTANCE_KEY_MTP3_CUTEDSL = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_cutedsl" - ACCEPTANCE_KEY_MTP5 = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp5" - ACCEPTANCE_KEY_MTP7 = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp7" - ACCEPTANCE_KEY_MTP7_LONG_CONTEXT = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp7_long_context" - ACCEPTANCE_KEY_MTP3_BF16_STATE = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_bf16_state" + ACCEPTANCE_KEY_MTP3_CUTEDSL_BF16_STATE = ( + "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_cutedsl_bf16_state" + ) ACCEPTANCE_KEY_MTP7_BF16_STATE_LONG_CONTEXT = ( "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp7_bf16_state_long_context" ) @@ -449,7 +442,6 @@ class TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1(LlmapiAccuracyTestHarness): ACCEPTANCE_KEY_NVFP4_KV_LONG_CONTEXT = ( "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_nvfp4_kv_long_context" ) - ACCEPTANCE_KEY_GDN_LEGACY = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp1_gdn_legacy" def _gsm8k_with_acceptance(self, mocker, monkeypatch, engine, speculative_config, key, label): """gsm8k against the reference row the engine resolves, then the acceptance length. @@ -484,62 +476,6 @@ def test_gsm8k(self, mocker, monkeypatch): task = GSM8K(self.MODEL_NAME) task.evaluate(llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) - @skip_not_sm103 - def test_gsm8k_breakable_prefill_cuda_graph(self, mocker, monkeypatch): - """The text model with ``prefill_cuda_graph_backend="breakable"``, against the same gsm8k anchor. - - Every prefill step of up to 3072 tokens is padded to a captured bucket and replayed, so - the padding rows, the four eager breaks and the per-bucket graphs run under real - prompts. That replays happen is shown by ``test_breakable_prefill_matches_eager``. - """ - mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", self.GSM8K_MAX_OUTPUT_LEN) - - with LLM( - self.MODEL_PATH, - **self.ENGINE_BREAKABLE, - **_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_gsm8k_cutedsl_moe(self, mocker, monkeypatch): - """The text model with the trunk's routed experts on CuteDSL, against the same gsm8k anchor. - - Same as ``test_gsm8k`` with ``moe_config.backend="CUTEDSL"``. - """ - mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", self.GSM8K_MAX_OUTPUT_LEN) - engine = dict(self.ENGINE, moe_config=MoeConfig(backend="CUTEDSL")) - - with LLM( - self.MODEL_PATH, - **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_gsm8k_bf16_ssm_state(self, mocker, monkeypatch): - """The text model with the recurrent state stored in bf16, against the fp32 gsm8k anchor. - - ``mamba_ssm_cache_dtype="bfloat16"`` (what ``auto`` resolves to for this checkpoint) rounds - the state to bf16 on every decode token and every prefill call, a lossy substitution - for the checkpoint's fp32 state. It reads the same reference row as ``test_gsm8k``, so - the rounding has to stay inside that row's tolerance. The decode step runs its - recurrence in the batch-invariant form that a bf16 pool needs. - """ - mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", self.GSM8K_MAX_OUTPUT_LEN) - engine = dict(self.ENGINE, kv_cache_config=self.KV_CACHE_BF16_STATE) - - with LLM( - self.MODEL_PATH, - **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. @@ -554,126 +490,28 @@ def test_mtp1_gsm8k_and_acceptance(self, mocker, monkeypatch): ) @skip_not_sm103 - def test_mtp1_fp8_kv_gsm8k_and_acceptance(self, mocker, monkeypatch): - """MTP with one draft over an fp8 KV cache: the FP8 KV row of references/gsm8k.yaml. - - ``kv_cache_config.dtype="fp8"`` makes the engine resolve ``kv_cache_quant_algo`` FP8, - which the reference lookup keys on: this case reads the FP8 row where - ``test_mtp1_gsm8k_and_acceptance`` (no KV quantization) reads the bf16 MTP row. The - draft layer runs over the same fp8 pool; its acceptance length has a floor of its own, - since fp8 storage can cost drafts the bf16 cache would have kept. - """ - engine = dict(self.ENGINE, kv_cache_config=self.KV_CACHE_FP8) - self._gsm8k_with_acceptance( - mocker, monkeypatch, engine, self.MTP1, self.ACCEPTANCE_KEY_FP8_KV, "mtp1 fp8 kv" - ) - - @skip_not_sm103 - def test_mtp1_nvfp4_kv_gsm8k_and_acceptance(self, mocker, monkeypatch): - """MTP with one draft over an NVFP4 KV cache: the NVFP4 KV row of references/gsm8k.yaml. - - ``kv_cache_config.dtype="nvfp4"`` makes the engine resolve ``kv_cache_quant_algo`` - NVFP4, which the reference lookup keys on. A lossy substitution of the checkpoint's bf16 - K/V (e2m1 values under e4m3 block scales), so the accuracy and the acceptance length - each have a row of their own. The draft layer runs over the same NVFP4 pool. - """ - engine = dict(self.ENGINE, kv_cache_config=self.KV_CACHE_NVFP4) - self._gsm8k_with_acceptance( - mocker, monkeypatch, engine, self.MTP1, self.ACCEPTANCE_KEY_NVFP4_KV, "mtp1 nvfp4 kv" - ) - - @skip_not_sm103 - def test_mtp3_gsm8k_and_acceptance(self, mocker, monkeypatch): - """MTP with three drafts: the MTP row of references/gsm8k.yaml, and a length floor. - - Three drafts replay the one MTP module three times, with the verify step running four - rows per request, a different shape from ``test_mtp1_gsm8k_and_acceptance`` on every - attention and recurrent-state call of the verify step. Drafts after the first can be - rejected while the first stays accepted, which holds the length near the one-draft - value, so the floor has its own entry. + def test_mtp3_cutedsl_bf16_ssm_state_gsm8k_and_acceptance(self, mocker, monkeypatch): + """Three drafts over a CuteDSL trunk and a bf16 recurrent state: the MTP row and a floor. + + The opt-in storage and kernel choices in one engine: the trunk's routed experts run on + CuteDSL next to the draft layer's FP8 block-scale experts, and the recurrent state is + bf16 (what ``auto`` resolves to for this checkpoint), rounded at every decode token, + prefill call and cached-replay commit. Three drafts give the verify step four rows per + request on every attention and recurrent-state call. Acceptance length detects a trunk + hidden state or a drifted state the draft layer cannot use, which accuracy does not. """ - self._gsm8k_with_acceptance( - mocker, monkeypatch, self.ENGINE, self.MTP3, self.ACCEPTANCE_KEY_MTP3, "mtp3" - ) - - @skip_not_sm103 - def test_mtp3_breakable_prefill_gsm8k_and_acceptance(self, mocker, monkeypatch): - """MTP with three drafts over replayed prefill steps: the MTP row, and the length floor. - - The padded rows must never reach the draft loop: the worker takes the batch's own rows - and the draft layer's first step runs over them eagerly. A draft layer fed a wrong - hidden state loses acceptance length without changing accuracy. - """ - self._gsm8k_with_acceptance( - mocker, - monkeypatch, - self.ENGINE_BREAKABLE, - self.MTP3, - self.ACCEPTANCE_KEY_MTP3, - "mtp3 breakable prefill", - ) - - @skip_not_sm103 - def test_mtp3_cutedsl_gsm8k_and_acceptance(self, mocker, monkeypatch): - """MTP with three drafts over a trunk whose routed experts run on CuteDSL. - - The trunk's CuteDSL experts and the draft layer's FP8 block-scale experts run in one - forward only here: ``test_gsm8k_cutedsl_moe`` has no draft layer and the other MTP - cases run the trunk on trtllm-gen. Acceptance length detects a trunk hidden state the - draft layer cannot use, which accuracy does not. - """ - engine = dict(self.ENGINE, moe_config=MoeConfig(backend="CUTEDSL")) - self._gsm8k_with_acceptance( - mocker, - monkeypatch, - engine, - self.MTP3, - self.ACCEPTANCE_KEY_MTP3_CUTEDSL, - "mtp3 cutedsl", - ) - - @skip_not_sm103 - def test_mtp5_gsm8k_and_acceptance(self, mocker, monkeypatch): - """MTP with five drafts: the MTP row of references/gsm8k.yaml, and a length floor. - - Six verify rows per request, a width that no other case runs: the cached replay at - ``T = 6`` and ``thop_attention`` at ``predicted_tokens_per_seq`` 6, with the all-layer - commit from 16 requests on. - """ - self._gsm8k_with_acceptance( - mocker, monkeypatch, self.ENGINE, self.MTP5, self.ACCEPTANCE_KEY_MTP5, "mtp5" - ) - - @skip_not_sm103 - def test_mtp7_gsm8k_and_acceptance(self, mocker, monkeypatch): - """MTP with seven drafts, the longest the target runs: the MTP row and a length floor. - - Eight verify rows per request is the cached replay kernel's ceiling (``T = 8``, history - 16): a slot that accepts every draft fills its history every other step, so the - checkpoint commit, in-kernel below 16 requests and all-layer from there on, runs twice - as often as at three drafts. The engine captures a CUDA graph for every batch size up - to 32, so every captured step carries up to 256 verify rows. - """ - self._gsm8k_with_acceptance( - mocker, monkeypatch, self.ENGINE, self.MTP7, self.ACCEPTANCE_KEY_MTP7, "mtp7" + engine = dict( + self.ENGINE, + moe_config=MoeConfig(backend="CUTEDSL"), + kv_cache_config=self.KV_CACHE_BF16_STATE, ) - - @skip_not_sm103 - def test_mtp3_bf16_ssm_state_gsm8k_and_acceptance(self, mocker, monkeypatch): - """MTP with three drafts over a bf16 recurrent state: the MTP row, and a length floor. - - The cached replay rebuilds the state from a checkpoint that is now bf16 and rounds it - once at every commit; drafts after the first are the ones a degraded state would lose, - so the length floor has its own entry. - """ - engine = dict(self.ENGINE, kv_cache_config=self.KV_CACHE_BF16_STATE) self._gsm8k_with_acceptance( mocker, monkeypatch, engine, self.MTP3, - self.ACCEPTANCE_KEY_MTP3_BF16_STATE, - "mtp3 bf16 state", + self.ACCEPTANCE_KEY_MTP3_CUTEDSL_BF16_STATE, + "mtp3 cutedsl bf16 state", ) def _long_context_passkey( @@ -730,11 +568,7 @@ def counting_reuse(*args, **kwargs): assert not missed, f"pass keys not stated back: {missed} of {[key for _, key in prompts]}" twins = outputs[len(prompts) // 2], outputs[-1] assert twins[0].outputs[0].token_ids == twins[1].outputs[0].token_ids - if engine.get("prefill_cuda_graph_backend") == "breakable": - # Every prompt is longer than the chunk budget, so its chunks fill the top bucket. - assert replays.get(2048, 0) >= len(texts), replays - else: - assert not replays, replays + assert not replays, replays if index_share: assert reuse_calls, "IndexShare is on and the reuse op never ran" else: @@ -746,28 +580,17 @@ def counting_reuse(*args, **kwargs): assert_acceptance_length(getattr(self, key_name), acceptance_length) @pytest.mark.parametrize( - "kv_cache, index_share, breakable, key_name", + "kv_cache, index_share, key_name", [ - pytest.param("KV_CACHE", True, False, "ACCEPTANCE_KEY_LONG_CONTEXT", id="index_share"), - pytest.param( - "KV_CACHE_FP8", False, False, "ACCEPTANCE_KEY_FP8_KV_LONG_CONTEXT", id="fp8_kv" - ), + pytest.param("KV_CACHE", True, "ACCEPTANCE_KEY_LONG_CONTEXT", id="index_share"), + pytest.param("KV_CACHE_FP8", False, "ACCEPTANCE_KEY_FP8_KV_LONG_CONTEXT", id="fp8_kv"), pytest.param( - "KV_CACHE_NVFP4", - False, - False, - "ACCEPTANCE_KEY_NVFP4_KV_LONG_CONTEXT", - id="nvfp4_kv", - ), - pytest.param( - "KV_CACHE", True, True, "ACCEPTANCE_KEY_LONG_CONTEXT", id="breakable_prefill" + "KV_CACHE_NVFP4", False, "ACCEPTANCE_KEY_NVFP4_KV_LONG_CONTEXT", id="nvfp4_kv" ), ], ) @skip_not_sm103 - def test_mtp3_long_context_passkey( - self, monkeypatch, kv_cache, index_share, breakable, key_name - ): + def test_mtp3_long_context_passkey(self, monkeypatch, kv_cache, index_share, key_name): """Pass key retrieval over 3k-6k-token prompts, where the sparse attention path runs. The gsm8k cases never leave the dense path (their prompts are under the indexer's @@ -779,15 +602,8 @@ def test_mtp3_long_context_passkey( ``index_share`` makes the draft loop reuse the selection of its first step; the acceptance length then must not fall below the floor taken from the same prompts with the draft layer selecting on every step, and the reuse op must have run. - - With ``breakable`` the prompts are prefilled in 2048-token chunks, each replayed from - the top bucket: the continuations past the first chunk read the KV, index and recurrent - state the earlier replays wrote, which the single-chunk gsm8k prompts never do. """ - engine = dict( - self.ENGINE_BREAKABLE_CHUNKED if breakable else self.ENGINE, - kv_cache_config=getattr(self, kv_cache), - ) + engine = dict(self.ENGINE, kv_cache_config=getattr(self, kv_cache)) self._long_context_passkey( monkeypatch, self.MTP3, @@ -795,7 +611,7 @@ def test_mtp3_long_context_passkey( index_share, _PASSKEY_LENGTHS, key_name, - f"mtp3 long context {kv_cache} breakable={breakable}", + f"mtp3 long context {kv_cache}", ) @skip_not_sm103 @@ -892,33 +708,17 @@ def run(engine): @skip_not_sm103 def test_mtp7_long_context_passkey(self, monkeypatch): - """Pass key retrieval over 3k-12k-token prompts with seven drafts and IndexShare on. - - Seven drafts run the sparse path with eight rows per request in the verify step and - draft decode steps that reuse the first step's selection up to six positions past it, - where the groups completed since the capture displace two of its ids. 12k tokens past + """Pass key retrieval over 3k-12k-token prompts: seven drafts, IndexShare, bf16 state. + + Seven drafts run the sparse path with eight rows per request in the verify step, the + cached replay kernel's ceiling (``T = 8``), and draft decode steps that reuse the first + step's selection up to six positions past it, where the groups completed since the + capture displace two of its ids. The recurrent state is bf16, what ``auto`` resolves to + for this checkpoint: it is rounded at every cached-replay checkpoint commit and once + per prefill chunk, so the rounding adds up with the length of the prompt and of the + answer, and pass keys stated back over 3k-12k tokens bound that drift. 12k tokens past the 8k default needs ``max_seq_len`` 16384. """ - engine = dict(self.ENGINE, max_seq_len=16384) - self._long_context_passkey( - monkeypatch, - self.MTP7, - engine, - True, - _PASSKEY_LENGTHS_LONG, - "ACCEPTANCE_KEY_MTP7_LONG_CONTEXT", - "mtp7 long context", - ) - - @skip_not_sm103 - def test_mtp7_bf16_ssm_state_long_context_passkey(self, monkeypatch): - """``test_mtp7_long_context_passkey`` over a bf16 recurrent state. - - The state is rounded to bf16 at every cached-replay checkpoint commit, each time a - slot's history of 16 accepted tokens fills, and once per prefill chunk, so the rounding - adds up with the length of the prompt and of the answer; pass keys stated back over - 3k-12k tokens bound that drift. - """ engine = dict(self.ENGINE, max_seq_len=16384, kv_cache_config=self.KV_CACHE_BF16_STATE) self._long_context_passkey( monkeypatch, @@ -932,24 +732,19 @@ def test_mtp7_bf16_ssm_state_long_context_passkey(self, monkeypatch): @pytest.mark.parametrize( "forced, accepted_rows", - [ - pytest.param("1e-9", 1, id="all_rejected"), - pytest.param("7", 8, id="all_accepted"), - ], + [pytest.param("1e-9", 1, id="all_rejected")], ) @skip_not_sm103 def test_mtp7_forced_acceptance_extremes(self, monkeypatch, forced, accepted_rows): - """Seven drafts with every draft rejected, then with every draft accepted, at batch 32. - - ``TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS`` overrides the sampler's acceptance: - a vanishing fraction accepts the target's own token and nothing else, ``7`` accepts - the seven drafts as well. 32 requests over 2k-6k-token prompts run the sparse path, - the all-layer replay commit and the CUDA graph at 32 x 8 = 256 verify rows with the - most rows rolled back at every step (seven) and with none rolled back, a history that - then fills to its end every second step. The acceptance length, read from the - requests' own metrics, must be exactly the forced row count: the acceptance - bookkeeping of the verify step moved by the number of rows forced, not by what the - drafts say. + """Seven drafts with every draft rejected, at batch 32. + + ``TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS`` overrides the sampler's acceptance: a + vanishing fraction accepts the target's own token and nothing else. 32 requests over + 2k-6k-token prompts run the sparse path, the all-layer replay commit and the CUDA graph + at 32 x 8 = 256 verify rows with the most rows rolled back at every step (seven). The + acceptance length, read from the requests' own metrics, must be exactly the forced row + count: the acceptance bookkeeping of the verify step moved by the number of rows + forced, not by what the drafts say. The text under forced acceptance is not compared with anything: the drafts a forced step accepts are not the target's tokens, and a run's tokens depend on the verify @@ -983,53 +778,6 @@ def test_mtp7_forced_acceptance_extremes(self, monkeypatch, forced, accepted_row print(f"[AL] mtp7 forced acceptance of {accepted_rows} rows: {acceptance_length:.3f}") assert acceptance_length == pytest.approx(accepted_rows, abs=0.05) - @skip_not_sm103 - def test_mtp1_gdn_legacy_verify(self, monkeypatch): - """MTP with one draft and ``TRTLLM_USE_GDN_REPLAY=0``: the per-step recurrent buffers. - - The engine runs the cached-replay verify step unless the variable is 0; this is the - only case that binds and runs the other one. A few word problems keep it short: their - answers say the recurrent state survives rejected drafts, and the acceptance length - says the drafts still land. - """ - monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") - tokenizer = AutoTokenizer.from_pretrained(self.MODEL_PATH) - prompts = [ - tokenizer.apply_chat_template( - [{"role": "user", "content": f"{question} {_WORD_PROBLEM_SUFFIX}"}], - tokenize=False, - add_generation_prompt=True, - enable_thinking=False, - ) - for question, _ in _WORD_PROBLEMS - ] - - with LLM( - self.MODEL_PATH, - speculative_config=self.MTP1, - **self.ENGINE, - **_require_target(monkeypatch, TRTLLM_USE_GDN_REPLAY="0", **_VALIDATE), - ) as llm: - outputs = llm.generate( - prompts, - SamplingParams(max_tokens=384, temperature=0, return_perf_metrics=True), - use_tqdm=False, - ) - bound = _target_core(llm)._bound - - assert bound.replay is None and bound.int_ssm, ( - "TRTLLM_USE_GDN_REPLAY=0 but the target bound the cached-replay verify state" - ) - wrong = [ - (question, answer, output.outputs[0].text[-120:]) - for (question, answer), output in zip(_WORD_PROBLEMS, outputs) - if f"answer is {answer}" not in output.outputs[0].text.replace("**", "") - ] - assert not wrong, f"wrong or missing answers: {wrong}" - acceptance_length = _acceptance_length_from_requests(outputs, self.MTP1.max_draft_len) - print(f"[AL] mtp1 gdn legacy acceptance_length = {acceptance_length:.3f}") - assert_acceptance_length(self.ACCEPTANCE_KEY_GDN_LEGACY, 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. @@ -1039,16 +787,6 @@ def test_mtp1_mmmu(self, monkeypatch): """ self._mmmu(monkeypatch, self.ENGINE) - @skip_not_sm103 - def test_mtp1_mmmu_breakable_prefill(self, monkeypatch): - """``test_mtp1_mmmu`` over replayed prefill steps, against the same row. - - The vision rows are spliced into the embedding inside the first break of every - replay; a replay that kept the previous step's rows or dropped them would answer - from the text alone, which this score detects. - """ - self._mmmu(monkeypatch, self.ENGINE_BREAKABLE) - def _mmmu(self, monkeypatch, engine): engine = dict(engine, kv_cache_config=self.KV_CACHE_MULTIMODAL) diff --git a/tests/integration/defs/accuracy/references/acceptance_length.yaml b/tests/integration/defs/accuracy/references/acceptance_length.yaml index aaa9f6e21d62..a87b9495bb69 100644 --- a/tests/integration/defs/accuracy/references/acceptance_length.yaml +++ b/tests/integration/defs/accuracy/references/acceptance_length.yaml @@ -154,43 +154,16 @@ ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp1: # measured, and tighten min_al then. min_al: 1.5 -# test_mtp1_fp8_kv_gsm8k_and_acceptance: the same draft over an fp8 KV cache, with a floor of -# its own. Storage in e4m3 can cost drafts the bf16 cache keeps, and the bf16 floor of 1.5 -# would let a quarter of the acceptance go unnoticed. -ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp1_fp8_kv: - # modeling_v2's own full-gsm8k reading on GB300 (TRTLLM MoE, 1319 samples); the bf16 KV - # reading of the same run is 1.971, so fp8 storage costs 0.4%. - ref_al: 1.964 - # 95% of ref_al. - min_al: 1.866 - -# test_mtp1_nvfp4_kv_gsm8k_and_acceptance: the same draft over an NVFP4 KV cache, with a floor of -# its own for the same reason as the fp8 one. -ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp1_nvfp4_kv: - # modeling_v2's own full-gsm8k reading on GB300 (TRTLLM MoE, 1319 samples); the bf16 KV - # reading of the same run is 1.971, so NVFP4 storage costs 0.6%. - ref_al: 1.959 - # 95% of ref_al. - min_al: 1.861 - -# test_mtp3_gsm8k_and_acceptance: three drafts, bf16 KV, trtllm-gen experts. -ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3: - # The built-in implementation's full-gsm8k reading on GB300, so the gate is not - # self-referential. modeling_v2's own is 3.704 (MTP2: 2.876). - ref_al: 3.567 - # Hand-set as a collapse tripwire, NOT the automatic 95%: 3.0 is well below both readings - # (16% under the built-in's, 19% under the target's) yet above what a broken later draft - # position scores. The first draft alone is the one-draft value, ~1.97; the first two - # without the third score ~2.8; a draft path rejected every time scores 1.0. Each of those - # leaves the text correct and every accuracy gate green. - min_al: 3.0 - -# test_mtp3_cutedsl_gsm8k_and_acceptance: three drafts with the trunk's experts on CuteDSL. -ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_cutedsl: - # modeling_v2's own full-gsm8k reading; there is no built-in reading for this combination - # (the built-in model has no CuteDSL trunk), so the floor is the mtp3 tripwire. - ref_al: 3.690 - min_al: 3.0 +# test_mtp3_cutedsl_bf16_ssm_state_gsm8k_and_acceptance: three drafts with the trunk's experts on +# CuteDSL over a bf16 recurrent state. +ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_cutedsl_bf16_state: + # modeling_v2's own full-gsm8k reading on GB300 (1319 samples); there is no built-in reading + # for this combination (the built-in model has no CuteDSL trunk). Each choice alone read + # 3.690 (CuteDSL) and 3.699 (bf16 state), against 3.704 for neither. + ref_al: 3.705 + # 95% of ref_al, above the two-draft reading (~2.8): a third draft position that is never + # accepted fails it. + min_al: 3.520 # test_mtp3_long_context_passkey[index_share]: three drafts with IndexShare on, bf16 KV, over # 13 prompts of 3k-6k tokens (the sparse attention path). Most of the output is a copy of the @@ -219,51 +192,11 @@ ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_nvfp4_kv_long_context: # 2.5% below ref_al, as for mtp3_long_context. min_al: 3.87 -# test_mtp1_gdn_legacy_verify: one draft with TRTLLM_USE_GDN_REPLAY=0, over six word problems. -ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp1_gdn_legacy: - # modeling_v2's own reading over those problems, 384 new tokens each; the cached-replay - # verify path reads the same over full gsm8k (1.971). - ref_al: 1.968 - # 95% of ref_al. - min_al: 1.87 - -# test_mtp5_gsm8k_and_acceptance: five drafts, bf16 KV, trtllm-gen experts. -ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp5: - # modeling_v2's own full-gsm8k reading on GB300 (1319 samples; four drafts read 4.428, six - # 5.609); there is no built-in reading at this draft length. - ref_al: 5.063 - # 95% of ref_al: above the four-draft reading, so a fifth draft position that is never - # accepted fails it. - min_al: 4.810 - -# test_mtp7_gsm8k_and_acceptance: seven drafts, the longest the target runs. -ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp7: - # modeling_v2's own full-gsm8k reading on GB300, as for mtp5. - ref_al: 6.052 - # 95% of ref_al: above the six-draft reading (5.609), so a seventh draft position that is - # never accepted fails it. - min_al: 5.749 - -# test_mtp7_long_context_passkey: seven drafts with IndexShare on over 12 prompts of 3k-12k -# tokens (the sparse attention path), whose output mostly copies the document's section -# headings: the readings sit near the range's ceiling of 8. -ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp7_long_context: - ref_al: 7.774 - # 2.5% below ref_al, as for mtp3_long_context. - min_al: 7.580 - -# test_mtp3_bf16_ssm_state_gsm8k_and_acceptance: three drafts over a bf16 recurrent state. -ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_bf16_state: - # modeling_v2's own full-gsm8k reading on GB300; the fp32-state reading is - # 3.704, so the bf16 state costs nothing measurable. - ref_al: 3.699 - # 95% of ref_al, above the two-draft reading (~2.8): a third draft position that is never - # accepted fails it. - min_al: 3.514 - -# test_mtp7_bf16_ssm_state_long_context_passkey: the mtp7_long_context prompts over a bf16 -# recurrent state (7.774 over fp32). +# test_mtp7_long_context_passkey: seven drafts with IndexShare on, over a bf16 recurrent state, +# over 12 prompts of 3k-12k tokens (the sparse attention path), whose output mostly copies the +# document's section headings: the readings sit near the range's ceiling of 8 (7.774 over an +# fp32 state). ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp7_bf16_state_long_context: ref_al: 7.857 - # 2.5% below ref_al, as for mtp7_long_context. + # 2.5% below ref_al, as for mtp3_long_context. min_al: 7.661 diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 4d0e2f9dc07a..480a51dc36cf 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -168,20 +168,17 @@ l0_b300: - unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k "TRTLLM" - unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k "CUTEDSL" - unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k "DEEPGEMM" - # qwen4_exp target, post-merge: one gate per opt-in variant -- three drafts, alone and - # over CuteDSL experts; seven drafts over 3k-12k-token pass keys (the verify step's - # eight rows per request); fp8 and NVFP4 KV over the pass keys (their sparse attention - # path); the recurrent state in bf16 at three drafts; seven drafts all rejected by - # forced acceptance (the speculative rollback end to end); synthetic videos; and the - # breakable prefill CUDA graph replayed against eager greedy continuations at the bucket - # edges and with context and generation rows in one batch. The other gates in - # test_qwen4_exp.py run on demand. - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_gsm8k_and_acceptance - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_cutedsl_gsm8k_and_acceptance - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp7_long_context_passkey + # qwen4_exp target, post-merge: the opt-in variants -- three drafts over CuteDSL experts and a + # bf16 recurrent state in one engine; fp8 and NVFP4 KV over the pass keys (their sparse + # attention path); seven drafts over 3k-12k-token pass keys on a bf16 state (the verify + # step's eight rows per request); seven drafts all rejected by forced acceptance (the + # speculative rollback end to end); synthetic videos; and the breakable prefill CUDA graph + # replayed against eager greedy continuations at the bucket edges and with context and + # generation rows in one batch. + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_cutedsl_bf16_ssm_state_gsm8k_and_acceptance - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_long_context_passkey[fp8_kv] - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_long_context_passkey[nvfp4_kv] - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_bf16_ssm_state_gsm8k_and_acceptance + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp7_long_context_passkey - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp7_forced_acceptance_extremes[all_rejected] - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp1_video - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_breakable_prefill_matches_eager From 078261b2e4974062d701b3502c5e1793e8c8b5b3 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Wed, 21 Oct 2026 02:10:55 -0700 Subject: [PATCH 24/28] [TRTLLM-17033][test] Qwen3.8-Flash-Next unit tests: certify the envelope the target runs The target's op tests drove every certified cell of every op: 2687 cases over 90 files. Each op's CELLS now covers the shapes the target calls it at plus one cell per distinct path the forward reaches (dtype, pad slots, layouts, kernel-switch boundaries), at most about six per entry; each note lists the paths that are no longer certified. Repeated fan-outs (graph replays, reject lists, per-cell composition checks, cold-process runs) keep one case per path. The two hc_decode_mix test files are merged, duplicated vision and pool-contract checks are dropped, and the CPU-only weight-table test is marked cpu_only and listed in l0_cpu.yml. The bf16 decode rows form is now certified at 128 rows, the size the target runs it at. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../activation/fused_sigmoid_mul_inplace.py | 73 +---- .../ops/activation/gelu_erf.py | 106 ++----- .../ops/activation/gelu_tanh.py | 111 ++----- .../ops/activation/silu_and_mul_rows.py | 1 - .../ops/attention/qsa_canonical_selection.py | 57 +--- .../ops/attention/qsa_expand_block_indices.py | 65 +--- .../ops/attention/qsa_mtp_reuse_topk.py | 114 +------ .../ops/attention/qsa_paged_index_scores.py | 290 +----------------- .../ops/attention/qsa_paged_index_store.py | 146 ++------- .../ops/attention/qsa_paged_kv_store.py | 132 ++------ .../ops/attention/qsa_paged_kv_store_nvfp4.py | 91 +----- .../ops/attention/qsa_paged_sparse_gqa.py | 151 +-------- .../attention/qsa_paged_sparse_gqa_decode.py | 4 +- .../attention/qsa_paged_sparse_gqa_nvfp4.py | 74 +---- .../ops/attention/qsa_prefill_compress.py | 67 +--- .../ops/attention/qsa_sparse_gqa_rowtile.py | 41 +-- .../ops/attention/qsa_topk_expand_decode.py | 12 - .../ops/attention/qsa_topk_select.py | 113 +------ .../qwen4_exp_attn_decode_prologue.py | 3 +- .../ops/attention/qwen4_exp_qk_norm_rope.py | 16 +- .../ops/attention/vision_rope_2d.py | 130 +++----- .../ops/embedding/ngram_hash.py | 30 +- .../ops/gemm/bf16_gemv_pdl.py | 23 +- .../ops/gemm/low_m_bf16_mm.py | 51 ++- .../ops/hc/qwen4_exp_gated_residual.py | 56 ++-- .../ops/hc/qwen4_exp_hc_decode_first_cute.py | 19 +- .../ops/hc/qwen4_exp_hc_decode_tc_first.py | 28 +- .../ops/hc/qwen4_exp_hc_inject_norm.py | 97 ++---- .../ops/hc/qwen4_exp_hc_lowrank_act.py | 7 +- .../ops/hc/qwen4_exp_hc_up_gate_mean.py | 14 +- .../ops/hc/qwen4_exp_hc_up_gate_mean_fp4.py | 11 +- .../ops/hc/qwen4_exp_post_inject.py | 67 ++-- ...sl_nvfp4_gather_grouped_gemm_act_fusion.py | 111 ++----- .../ops/moe/cute_dsl_nvfp4_grouped_gemm.py | 124 +++----- .../fp4_block_scale_moe_renorm_permuted.py | 10 +- .../moe/fp4_block_scale_moe_renorm_routed.py | 14 +- .../ops/moe/fp4_moe_renorm_cute_dsl.py | 15 +- .../ops/moe/fp8_block_scale_moe_runner.py | 82 +---- .../ops/moe/fused_sigmoid_gate_mul_add.py | 34 +- .../ops/moe/moe_sort.py | 86 ++---- .../ops/moe/moe_unpermute_inplace.py | 28 +- .../ops/moe/renorm_moe_routing_op.py | 33 +- .../ops/norm/flashinfer_gemma_rmsnorm.py | 47 +-- .../ops/norm/layer_norm.py | 61 +--- .../ops/norm/rms_norm_gated_token_major.py | 61 +--- .../ops/ple/gated_ngram_injection.py | 71 ++--- .../ops/ple/ngram_token_history.py | 214 +++---------- .../ops/ple/ple_gate_conv_inject_norm.py | 35 +-- .../ops/ple/ple_ngram_embed.py | 111 +++---- .../ops/quantization/bf16_to_fp32.py | 5 +- .../dequantize_e4m3_per_tensor.py | 67 +--- .../ops/quantization/fp8_quantize_1x128.py | 35 +-- .../ops/ssm/causal_conv1d_fwd.py | 227 ++++---------- .../ops/ssm/causal_conv1d_triton_update.py | 169 ++-------- .../ops/ssm/causal_conv1d_update.py | 177 ++--------- .../ops/ssm/chunk_gated_delta_rule.py | 155 ++++------ .../ops/ssm/cu_seqlens_from_seq_lens.py | 88 ++---- .../ops/ssm/dilated_short_conv_paged.py | 177 +++-------- .../ssm/extract_transpose_prefill_slice.py | 129 ++------ .../ssm/flashinfer_chunk_gated_delta_rule.py | 37 +-- .../ops/ssm/fused_gdn_gating.py | 110 ++----- .../ops/ssm/fused_gdn_post_conv.py | 84 ++--- ...t_gated_delta_rule_cached_replay_update.py | 160 +--------- ...fused_recurrent_gated_delta_rule_update.py | 141 ++------- .../fused_sigmoid_gating_delta_rule_update.py | 214 +++---------- .../ops/ssm/gdn_in_proj_conv1d_pdl.py | 1 - .../ops/ssm/gdn_prefill_conv_post.py | 60 +--- .../ops/ssm/pack_gdn_decode_qkv.py | 75 +---- .../ops/ssm/reset_recurrent_state_rows.py | 101 +----- .../ops/vision/patch_embed_conv3d.py | 74 +---- .../ops/vision/pos_embed_interpolate.py | 90 ++---- .../integration/test_lists/test-db/l0_cpu.yml | 3 +- ...t_modeling_v2_fused_sigmoid_mul_inplace.py | 54 +--- .../activation/test_modeling_v2_gelu_erf.py | 20 +- .../activation/test_modeling_v2_gelu_tanh.py | 20 +- .../test_modeling_v2_silu_and_mul_rows.py | 2 +- ...est_modeling_v2_qsa_canonical_selection.py | 21 +- ...st_modeling_v2_qsa_expand_block_indices.py | 21 +- .../test_modeling_v2_qsa_mtp_reuse_topk.py | 17 +- ...test_modeling_v2_qsa_paged_index_scores.py | 101 ++---- .../test_modeling_v2_qsa_paged_index_store.py | 120 ++------ .../test_modeling_v2_qsa_paged_kv_store.py | 8 +- ...st_modeling_v2_qsa_paged_kv_store_nvfp4.py | 27 +- .../test_modeling_v2_qsa_paged_sparse_gqa.py | 26 +- ...modeling_v2_qsa_paged_sparse_gqa_decode.py | 8 +- ..._modeling_v2_qsa_paged_sparse_gqa_nvfp4.py | 12 +- .../test_modeling_v2_qsa_prefill_compress.py | 34 +- ...test_modeling_v2_qsa_sparse_gqa_rowtile.py | 17 +- ...test_modeling_v2_qsa_topk_expand_decode.py | 46 +-- .../test_modeling_v2_qsa_topk_select.py | 33 +- ...test_modeling_v2_qwen4_exp_qk_norm_rope.py | 6 +- .../embedding/test_modeling_v2_ngram_hash.py | 11 +- .../gemm/test_modeling_v2_bf16_gemv_pdl.py | 10 +- .../gemm/test_modeling_v2_low_m_bf16_mm.py | 52 +++- ...eling_v2_qwen4_exp_hc_decode_first_cute.py | 8 +- ...odeling_v2_qwen4_exp_hc_decode_mix_cute.py | 118 ------- ...ing_v2_qwen4_exp_hc_decode_mix_fp4_cute.py | 56 ++-- ...odeling_v2_qwen4_exp_hc_decode_tc_first.py | 14 +- ...st_modeling_v2_qwen4_exp_hc_inject_norm.py | 38 +-- ...t_modeling_v2_qwen4_exp_hc_up_gate_mean.py | 18 +- ...deling_v2_qwen4_exp_hc_up_gate_mean_fp4.py | 16 +- ...sl_nvfp4_gather_grouped_gemm_act_fusion.py | 73 ++--- ...modeling_v2_cute_dsl_nvfp4_grouped_gemm.py | 77 ++--- .../test_modeling_v2_cute_dsl_weight_path.py | 15 +- ...ng_v2_fp4_block_scale_moe_renorm_routed.py | 2 +- ...est_modeling_v2_fp4_moe_renorm_cute_dsl.py | 12 +- ..._modeling_v2_fp8_block_scale_moe_runner.py | 5 +- ..._modeling_v2_fused_sigmoid_gate_mul_add.py | 33 +- .../moe/test_modeling_v2_moe_sort.py | 32 +- .../test_modeling_v2_moe_unpermute_inplace.py | 32 +- ...st_modeling_v2_flashinfer_gemma_rmsnorm.py | 7 +- .../norm/test_modeling_v2_layer_norm.py | 20 +- ..._modeling_v2_rms_norm_gated_token_major.py | 54 ++-- .../test_modeling_v2_gated_ngram_injection.py | 8 +- .../test_modeling_v2_ngram_token_history.py | 37 +-- ...t_modeling_v2_ple_gate_conv_inject_norm.py | 15 +- .../ple/test_modeling_v2_ple_ngram_embed.py | 25 +- ..._modeling_v2_dequantize_e4m3_per_tensor.py | 16 +- .../test_modeling_v2_fp8_quantize_1x128.py | 4 +- .../ssm/test_modeling_v2_causal_conv1d_fwd.py | 42 +-- ...modeling_v2_causal_conv1d_triton_update.py | 15 +- .../test_modeling_v2_causal_conv1d_update.py | 45 +-- ...test_modeling_v2_chunk_gated_delta_rule.py | 58 +--- ...st_modeling_v2_cu_seqlens_from_seq_lens.py | 14 +- ...st_modeling_v2_dilated_short_conv_paged.py | 15 +- ...ling_v2_extract_transpose_prefill_slice.py | 10 +- .../ssm/test_modeling_v2_fused_gdn_gating.py | 51 +-- .../test_modeling_v2_fused_gdn_post_conv.py | 16 +- ...t_gated_delta_rule_cached_replay_update.py | 25 +- ...fused_recurrent_gated_delta_rule_update.py | 6 +- ..._fused_sigmoid_gating_delta_rule_update.py | 18 +- .../ssm/test_modeling_v2_gdn_block.py | 12 +- ...t_modeling_v2_gdn_decode_gated_norm_pdl.py | 4 +- ..._modeling_v2_gdn_decode_rows_gated_norm.py | 15 +- ...test_modeling_v2_gdn_in_proj_conv1d_pdl.py | 12 +- .../test_modeling_v2_gdn_prefill_conv_post.py | 52 ++-- .../test_modeling_v2_pack_gdn_decode_qkv.py | 74 ++--- ..._modeling_v2_reset_recurrent_state_rows.py | 46 +-- .../test_modeling_v2_gdn_replay_verify.py | 11 +- .../test_modeling_v2_index_share.py | 56 +--- .../test_modeling_v2_kv_range_probe.py | 2 +- .../test_modeling_v2_logits_widening.py | 2 +- .../test_modeling_v2_spec_rollback.py | 8 +- .../test_modeling_v2_patch_embed_conv3d.py | 39 +-- .../test_modeling_v2_pos_embed_interpolate.py | 7 +- .../test_modeling_v2_vision_tower_calls.py | 44 +-- .../test_modeling_v2_qwen4_exp_kv_quant.py | 11 +- .../test_modeling_v2_qwen4_exp_moe_scalars.py | 2 +- ...modeling_v2_qwen4_exp_multimodal_embeds.py | 6 +- ...est_modeling_v2_qwen4_exp_pool_contract.py | 18 +- ..._modeling_v2_qwen4_exp_validate_request.py | 6 +- ...test_modeling_v2_qwen4_exp_weight_table.py | 3 +- 152 files changed, 1823 insertions(+), 6039 deletions(-) delete mode 100644 tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_cute.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 index b870dd1627c1..deee31706078 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/fused_sigmoid_mul_inplace.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/fused_sigmoid_mul_inplace.py @@ -34,15 +34,11 @@ class FusedSigmoidMulInplace(OpWrapper): CELLS: tuple[Cell, ...] = ( Cell( - why="decode: 1 token, 24x256 gate as the strided half of q_proj", + why="decode: 1 token, 24x256 gate as the strided half of q_proj (MTP)", 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", + why="prefill 2048 tokens, strided q_proj half: >= 1024 rows picks 2048-wide blocks", spec=dict(tokens=2048, **_CKPT), ), Cell( @@ -57,53 +53,6 @@ class FusedSigmoidMulInplace(OpWrapper): why="prefill 8192 tokens: 2-D gate off the fused projection row", spec=dict(tokens=8192, heads=24, head_dim=256, dtype=torch.bfloat16, gate="fused_row"), ), - 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), @@ -122,22 +71,26 @@ class FusedSigmoidMulInplace(OpWrapper): `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). + The pre-call contents of `attention_output` are destroyed; the gate is + left bitwise unchanged. - Numerics: bf16/fp16 bitwise equal to the fp32 reference except ~1e-6 of + Numerics: bf16 bitwise equal to the fp32 reference except ~1e-6 of elements one ulp off (`tl.sigmoid`'s fast exp); default tolerance holds, nothing loosened. Feeding the q half instead of the gate half moves the result by ~2.15 max abs at [17, 6144], so the per-head indexing is under test. + Certified: bf16 operands at the checkpoint's hidden 6144, a compact + `attention_output`, and the target's two gate layouts (the strided q_proj + half and the 2-D window of the fused projection row). + Hard limits T * row stride token offset computed in int64 boundary 2^63 not probed - (largest run 2048 rows) + (largest run 8192 rows) - Not certified: CUDA-graph capture, torch.compile. + Not certified: a strided `attention_output`, a gate dtype other than + `attention_output.dtype`, fp16/fp32 operands, other hidden sizes, + CUDA-graph capture, torch.compile. """ def __call__(self, attention_output: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_erf.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_erf.py index 8870a7a5eeea..edae62e6738d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_erf.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_erf.py @@ -23,11 +23,10 @@ #: `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 -- and 8192 the most one tower call carries (32768 patch -#: rows). -MERGER_ROW_COUNTS = (1, 7, 64, 65, 140, 1024, 4096, 8192) +#: Rows the merger is asked for, `sum(t*h*w) / 4`: 1 the smallest, 140 a +#: 336x448 image (grid [1, 20, 28]) -- not a multiple of any MoE tile size -- +#: and 8192 the most one tower call carries (32768 patch rows). +MERGER_ROW_COUNTS = (1, 140, 8192) _SQRT_2 = math.sqrt(2.0) @@ -46,90 +45,43 @@ class GeluErf(OpWrapper): # spec: shape, dtype; `dist` "uniform" over [lo, hi] (default) or "merger" # (a tower hidden state scaled by the merger norm's measured gamma, grouped - # 4-to-1, through an fc1 at the checkpoint's measured weight/bias scale); - # `layout` "contiguous" (default), "column_slice" or "transposed". - 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), - ), + # 4-to-1, through an fc1 at the checkpoint's measured weight/bias scale). + CELLS: tuple[Cell, ...] = tuple( + Cell( + why=f"merger activand, {rows} merged rows x 4608 bf16", + spec=dict(shape=(rows, MERGER_WIDTH), dtype=torch.bfloat16, lo=-16.0, hi=16.0), ) + for rows in MERGER_ROW_COUNTS + ) + ( + Cell( + why="activand built the way the merger builds it", + spec=dict(shape=(512, MERGER_WIDTH), dtype=torch.bfloat16, dist="merger"), + ), + Cell( + why="float32 over [-3, 3], the erf/tanh divergence peak", + spec=dict(shape=(256, MERGER_WIDTH), dtype=torch.float32, lo=-3.0, hi=3.0), + ), ) note = """ Metadata consumed: none. Out-of-place; `x` is not mutated. - The output inherits `x`'s memory format: a transposed `x` returns a - non-contiguous result (bitwise equal to the call on a contiguous copy). + Certified on contiguous bf16 `[rows, 4608]` activands (what the merger's fc1 + returns) plus one fp32 cell. - Identity is pinned by the fp32 cells only. Max |erf GELU - tanh GELU| is - 4.73e-4 at x = -2.699; in fp32 the tanh form lands ~1400x the correct path's - deviation away, 31x outside the default gate. In bf16 one ulp at the output - scale exceeds that gap (tanh variant 8.24e-3 vs correct 7.77e-3 max abs), so - a bf16 match does not say which GELU was computed. + Identity is pinned by fp32 only (the fp32 cell and the identity test). + Max |erf GELU - tanh GELU| is 4.73e-4 at x = -2.699; in fp32 the tanh form + lands ~1400x the correct path's deviation away, 31x outside the default + gate. In bf16 one ulp at the output scale exceeds that gap (tanh variant + 8.24e-3 vs correct 7.77e-3 max abs), so a bf16 match does not say which GELU + was computed. Values: `-inf -> nan` (the closed form does the same); `+inf -> +inf`; large negatives give `-0.0`; `finfo.max` comes back unchanged. CUDA-graph capturable; the output is a fresh allocation. - Not certified: float8 operands. No index-typed argument; numel past 2^31 + Not certified: non-contiguous inputs, fp16, other widths, float8 + operands. No index-typed argument; numel past 2^31 not probed. """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_tanh.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_tanh.py index a4ef7a0fe27a..53db4ce79ebb 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_tanh.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_tanh.py @@ -40,101 +40,48 @@ class GeluTanh(OpWrapper): # spec: shape, dtype; `dist` "uniform" over [lo, hi] (default) or "fc1" (a # unit-normal hidden state through an fc1 at the checkpoint's measured - # weight std 0.0197 / bias std 1.676); `layout` "contiguous" (default), - # "column_slice" (the first half of a twice-as-wide buffer) or "transposed". - 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="block MLP activand at the tower call ceiling, 32768 patch rows x 4304 bf16", - spec=dict( - shape=(32768, INTERMEDIATE_SIZE), dtype=torch.bfloat16, lo=-16.0, hi=16.0 - ), - ), - ) - + ( - Cell( - why="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), - ), + # weight std 0.0197 / bias std 1.676). + CELLS: tuple[Cell, ...] = tuple( + Cell( + why=f"block MLP activand, {rows} patch rows x 4304 bf16", + spec=dict(shape=(rows, INTERMEDIATE_SIZE), dtype=torch.bfloat16, lo=-16.0, hi=16.0), ) + for rows in (1, 1024) + ) + ( + Cell( + why="block MLP activand at the tower call ceiling, 32768 patch rows x 4304 bf16", + spec=dict(shape=(32768, INTERMEDIATE_SIZE), dtype=torch.bfloat16, lo=-16.0, hi=16.0), + ), + Cell( + why="activand built through fc1 at the checkpoint's measured scale", + spec=dict(shape=(512, INTERMEDIATE_SIZE), dtype=torch.bfloat16, dist="fc1"), + ), + Cell( + why="float32 over [-3, 3], the tanh/erf divergence peak", + spec=dict(shape=(256, INTERMEDIATE_SIZE), dtype=torch.float32, lo=-3.0, hi=3.0), + ), ) note = """ Metadata consumed: none. Out-of-place; `x` is not mutated. - The output inherits `x`'s memory format: a transposed `x` returns a - non-contiguous result (bitwise equal to the call on a contiguous copy). + Certified on contiguous bf16 `[rows, 4304]` activands (what the block MLP's + fc1 returns) plus one fp32 cell. - Identity is pinned by the fp32 cells only. Max |tanh GELU - erf GELU| is - 4.73e-4 (|x| ~ 2.7); in fp32 the erf form lands ~1600x the correct path's - deviation away, 31x outside the default gate. In bf16 one ulp at the output - scale exceeds that gap (erf variant 8.19e-3 vs correct 7.75e-3 max abs), so a - bf16 match does not say which GELU was computed. + Identity is pinned by fp32 only (the fp32 cell and the identity test). + Max |tanh GELU - erf GELU| is 4.73e-4 (|x| ~ 2.7); in fp32 the erf form + lands ~1600x the correct path's deviation away, 31x outside the default + gate. In bf16 one ulp at the output scale exceeds that gap (erf variant + 8.19e-3 vs correct 7.75e-3 max abs), so a bf16 match does not say which GELU + was computed. Values: `-inf -> nan` (the closed form does the same: -inf times a factor saturated to 0); `+inf -> +inf`; large negatives give `-0.0`; `finfo.max` comes back unchanged in bf16/fp16/fp32 (the cube is not taken in the input dtype). CUDA-graph capturable; the output is a fresh allocation. - Not certified: float8 operands. No index-typed argument; numel past 2^31 + Not certified: non-contiguous inputs, fp16, other widths, float8 + operands. No index-typed argument; numel past 2^31 not probed. """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/silu_and_mul_rows.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/silu_and_mul_rows.py index ad0d9457514d..a79777be951b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/silu_and_mul_rows.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/silu_and_mul_rows.py @@ -50,7 +50,6 @@ class SiluAndMulRows(OpWrapper): Cell(why="prefill chunk 16384 rows of the fused 1288-wide output", spec=dict(rows=16384)), Cell(why="decode batch 128, fused output", spec=dict(rows=128)), Cell(why="1 row, fused output", spec=dict(rows=1)), - Cell(why="33 rows, contiguous gate|up (no extra columns)", spec=dict(rows=33, width=1280)), Cell(why="zero rows", spec=dict(rows=0)), ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_canonical_selection.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_canonical_selection.py index 130518c7111c..ddd0b52d6281 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_canonical_selection.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_canonical_selection.py @@ -103,19 +103,16 @@ class QsaCanonicalSelection(OpWrapper): # value, in a random order (the adversary); `radix` is the engine the target calls. CELLS: tuple[Cell, ...] = ( _cell( - "decode: 1, 8, 128 and 256 rows (MTP3 at batch 32, MTP7 at batch 32)", - (1, 8, 128, 256), + "decode: 1, 128 and 256 rows (MTP3 and MTP7 at batch 32)", + (1, 128, 256), ("positions_random", 2051, 8191), "relu", ), - _cell("prefill: 1024 rows from position 4000", 1024, ("chunk", 4000), "relu"), - _cell("prefill: 8192 rows from position 0", 8192, ("chunk", 0), "relu"), _cell( - "prefill: 8192 rows over an 8192-wide table, rewritten in row tiles", + "prefill: 8192 rows from position 0, rewritten in row tiles", 8192, - ("chunk", 24000), - "relu_sparse", - cols=8192, + ("chunk", 0), + "relu", ), _cell( "windows 0, 1, 511, 512, 513, 1024, 2047, 2048, 2080", @@ -124,53 +121,11 @@ class QsaCanonicalSelection(OpWrapper): "relu", ), _cell( - "ties: fewer than k positive scores, the zero class straddles k", - 64, - ("ends", (2048,) * 32 + (1500,) * 16 + (2080,) * 16), - "relu_sparse", - ), - _cell( - "ties: two-valued rows", - 32, - ("ends", (2048,) * 16 + (777,) * 16), - "two_valued", - ), - _cell( - "ties: every window column equal", - 16, - ("ends", (2048,) * 8 + (513,) * 8), - "all_equal", - ), - _cell( - "ties: one positive pair tied exactly at the k-th value", - 16, - ("ends", (2048,) * 16), - "pair_at_k", - ), - _cell( - "ties: -0.0 and +0.0 mixed in the zero class", + "ties: a zero class of -0.0 and +0.0 that straddles k", 16, ("ends", (2048,) * 16), "signed_zeros", ), - _cell( - "tail of +inf / 3e38 / NaN past the window is never kept", - 32, - ("positions_random", 2051, 8191), - "relu_sparse", - tail="poison", - ), - _cell( - "nonzero row_starts: offsets stay window-local", - 8, - ( - "explicit", - (1, 2, 3, 37, 100, 480, 1001, 1500), - (2080, 600, 515, 1061, 612, 2079, 1513, 2080), - ), - "relu_sparse", - tail="poison", - ), _cell( "input from the radix engine (qsa_topk_select), zero ties", 64, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_expand_block_indices.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_expand_block_indices.py index 0744a4ca1f2d..bacb567ea3e0 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_expand_block_indices.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_expand_block_indices.py @@ -42,10 +42,6 @@ def _spec(batch, **overrides) -> dict: 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. @@ -76,44 +72,16 @@ class QsaExpandBlockIndices(OpWrapper): ARCHS = frozenset({Arch.SM_103}) CELLS: tuple[Cell, ...] = ( - # ---- decode rows: the column-split launch below 128 rows ------------- + # ---- decode rows: the column-split launch below 128 rows, whole rows from 128 Cell(why="decode: one row, columns split across programs", spec=_spec(("sweep", 1))), - Cell(why="decode MTP: one request's two rows", spec=_spec(("sweep", 2))), - Cell(why="decode: one row for each of 32 request slots", spec=_spec(("sweep", 32))), - Cell(why="127 rows: the last row count on the column split", spec=_spec(("sweep", 127))), Cell(why="128 rows: the first row count on whole-row programs", spec=_spec(("sweep", 128))), - Cell(why="256 rows on whole-row programs", spec=_spec(("sweep", 256))), # ---- MTP verify blocks (max_draft_len up to 7): p rows per request, every # row's length the block's end, so only the last row has length == position + 1. - Cell( - why="verify MTP7: one request's 8 rows, column split", - spec=_spec(("verify", (6000,), 8)), - ), - Cell( - why="verify MTP7: 15 requests x 8 = 120 rows, column split", - spec=_spec(("verify", _VERIFY_STARTS[:15], 8)), - ), - Cell( - why="verify MTP7: 16 requests x 8 = 128 rows, whole-row programs", - spec=_spec(("verify", _VERIFY_STARTS[:16], 8)), - ), Cell( why="verify MTP7: 32 requests x 8 = 256 rows, whole-row programs", spec=_spec(("verify", _VERIFY_STARTS, 8)), ), - # ---- 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,))), - ), + # ---- context rows: sequence_lengths > query_position + 1 inside chunks Cell( why="warm-up capture: 2 x 4096 rows, the [8192, 512] buffer's height", spec=_spec(("serving", 2, (4096, 4096))), @@ -122,39 +90,10 @@ class QsaExpandBlockIndices(OpWrapper): 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 = """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_mtp_reuse_topk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_mtp_reuse_topk.py index ad419f0cd502..89367af18c42 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_mtp_reuse_topk.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_mtp_reuse_topk.py @@ -90,15 +90,7 @@ class QsaMtpReuseTopk(OpWrapper): ARCHS = frozenset({Arch.SM_103}) CELLS: tuple[Cell, ...] = ( - # ---- single requests: the two behaviours --------------------------- - Cell( - why="1 request, short row: group 100 appended after a 100-id prefix", - spec=_spec(("explicit", (402,)), drafted=(1,), shared_rows=8), - ), - Cell( - why="1 request, full row: group 512 displaces the oldest id, slot 0", - spec=_spec(("explicit", (2050,)), drafted=(1,), shared_rows=8), - ), + # ---- the two behaviours, at every boundary ------------------------- Cell( why="8 requests, one per boundary: 0/1/100/511/512/2047 groups, fresh 0 and 1", spec=_spec( @@ -107,41 +99,6 @@ class QsaMtpReuseTopk(OpWrapper): shared_rows=8, ), ), - # ---- decode batches: one row per request ---------------------------- - Cell(why="decode MTP: 4 requests, mixed rows", spec=_spec(("mixed", 4))), - Cell( - why="decode MTP: 32 requests (envelope), 256-row shared buffer", - spec=_spec(("mixed", 32), shared_rows=256), - ), - Cell( - why="decode MTP: 64 requests, the shared buffer exactly full", - spec=_spec(("mixed", 64), shared_rows=64), - ), - # ---- positions drafted since the capture: draft step i is i past it -- - Cell( - why="32 requests, every one 1 position past its capture", - spec=_spec(("mixed", 32), drafted=(1,)), - ), - Cell( - why="32 requests, every one 2 positions past its capture", - spec=_spec(("mixed", 32), drafted=(2,)), - ), - Cell( - why="32 requests, every one 3 positions past its capture", - spec=_spec(("mixed", 32), drafted=(3,)), - ), - Cell( - why="32 requests, every one 4 positions past its capture", - spec=_spec(("mixed", 32), drafted=(4,)), - ), - Cell( - why="32 requests, every one 5 positions past its capture", - spec=_spec(("mixed", 32), drafted=(5,)), - ), - Cell( - why="32 requests, every one 6 positions past its capture (max_draft_len 7)", - spec=_spec(("mixed", 32), drafted=(6,)), - ), Cell( why="8 requests at fresh 2: full, 511-id and short rows gain two groups", spec=_spec( @@ -150,77 +107,20 @@ class QsaMtpReuseTopk(OpWrapper): shared_rows=8, ), ), - # ---- captured-row fill --------------------------------------------- - Cell( - why="32 short contexts: every captured row below 512 ids", - spec=_spec(("short", 32)), - ), - Cell( - why="32 contexts past 2048 tokens: every captured row full", - spec=_spec(("full", 32)), - ), - Cell( - why="32 short contexts, 4..6 positions past the capture: fresh up to 2", - spec=_spec(("short", 32), drafted=(4, 5, 6)), - ), + # ---- decode batches: one row per request ---------------------------- Cell( - why="32 full rows, 4..6 positions past the capture: up to 2 oldest ids displaced", - spec=_spec(("full", 32), drafted=(4, 5, 6)), + why="decode MTP: 32 requests (envelope), 256-row shared buffer", + spec=_spec(("mixed", 32), shared_rows=256), ), - # ---- addressing: the request map and the slot ----------------------- Cell( why="permuted request map, draft manager's single slot", spec=_spec(("mixed", 32), request_map="permuted", slots=_DRAFT_SLOTS, shared_rows=32), ), - Cell( - why="3 rows per request (24 rows, 8 requests) share one captured row", - spec=_spec(("mixed", 8), rows_per_request=3, shared_rows=8), - ), - Cell( - why="group gaps 2..600 and current < captured (not reached at cr 4)", - spec=_spec( - ( - "gaps", - ( - (0, 600), - (10, 15), - (500, 1100), - (100, 98), - (2047, 2048), - (300, 812), - (300, 811), - ), - ), - shared_rows=8, - ), - ), # ---- CUDA graph ----------------------------------------------------- - Cell( - why="CUDA graph: captured at 1, 8 and 32 requests, new-data replay", - spec=_spec(("mixed", 64), graph_batches=(1, 8, 32)), - ), Cell( why="CUDA graph, drafted 1..6: captured at 1, 8 and 32 requests, new-data replay", spec=_spec(("mixed", 32), drafted=(1, 2, 3, 4, 5, 6), graph_batches=(1, 8, 32)), ), - # ---- dtypes, layouts, geometry -------------------------------------- - Cell( - why="int64 request map and visible counts", - spec=_spec(("mixed", 8), metadata_dtype=torch.int64, shared_rows=8), - ), - Cell( - why="dense standalone tensors instead of engine buffer slices", - spec=_spec(("mixed", 8), layout="dense", shared_rows=8), - ), - Cell( - why="geometry block_topk 8", - spec=_spec(("explicit", (2, 30, 34, 3, 41, 13)), block_topk=8, shared_rows=8), - ), - Cell( - why="geometry block_topk 1024", - spec=_spec(("explicit", (2, 4094, 4098, 402, 8188)), block_topk=1024, shared_rows=8), - ), - Cell(why="zero rows: output returned, no launch", spec=_spec(("mixed", 0))), ) note = """ @@ -239,7 +139,7 @@ class QsaMtpReuseTopk(OpWrapper): `max_draft_len` 1 never reaches this op, 3 reaches 1 and 2 positions, and 7 (the target's ceiling) reaches 1..6. With compress_ratio 4, `fresh = ((p + 1) % 4 + i) // 4` for a capture at position `p`: 0 or 1 up to i = 3 - and 0, 1 or 2 from i = 4; the gap cell covers wider gaps. In a full row + and 0, 1 or 2 from i = 4. In a full row the fresh ids displace the row's oldest ids (one at fresh 1, two at fresh 2), and a 511-id row at fresh 2 loses its oldest id. @@ -265,9 +165,9 @@ class QsaMtpReuseTopk(OpWrapper): a 2-D `shared_visible_blocks` (the whole `[slots, seqs]` buffer is read as slot 0's counts). - Certified: int32 ids, int32 or int64 metadata, `block_topk` 8, 512 and 1024. + Certified: int32 ids and metadata, `block_topk` 512, up to 32 requests. No host synchronization and shape-only validation, so it captures into a - CUDA graph; `rows == 0` returns `output` without a launch. + CUDA graph. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_scores.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_scores.py index 6675d234850f..fc41479f5980 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_scores.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_scores.py @@ -35,41 +35,10 @@ def _spec(**overrides) -> dict: return {**_CHECKPOINT, **overrides} -def _small(rows: int, requests: int, **overrides) -> dict: - """The narrow-table cells: random positions over an 8-page table, rows - grouped by request, each request's table row fully mapped, a contiguous q. - Eight table rows, as in the serving dump, whatever the request count.""" - base = dict(layout="random", rows=rows, requests=requests, slots=8, pages=8) - base.update(q_layout="packed", positions_dtype=torch.int32) - base.update(only_visible_blocks=False, context_rows=False) - base["branch"] = "row_wise" - base.update(overrides) - return _spec(**base) - - #: The serving runtime dumps: an [8, 260] table on the warm-up capture path and #: a [32, 260] one under the mmmu gate's mixed steps. _WIDTH = 260 -_FREE_HEADS = tuple( - Cell( - why=f"free axis, 64 rows: {h} index heads x dim {d}", - spec=_small(64, 3, heads=h, head_dim=d), - ) - for d in (64, 128) - for h in (1, 2, 3, 4, 8, 16) - if (h, d) != (4, 128) -) -_FREE_PAGING = tuple( - Cell( - why=f"free axis, 48 rows: compress {c} x {t} tokens per block", - spec=_small(48, 3, compress_ratio=c, tokens_per_block=t), - ) - for t in (16, 32, 64) - for c in (2, 4, 8, 16) - if (c, t) != (4, 32) -) - #: The same plane under an NVFP4 KV pool: the QSA position page joins each physical #: page after the index buffers (`page_pad` elements), so the trunk's view is @@ -112,7 +81,6 @@ class QsaPagedIndexScores(OpWrapper): ARCHS = frozenset({Arch.SM_103}) CELLS: tuple[Cell, ...] = ( - # ---- the NVFP4-pool plane: page stride 49344, a kernel of its own ---- Cell( why="NVFP4 trunk plane, page stride 49344: decode MTP 4 requests x 2 rows", spec=_spec( @@ -140,7 +108,6 @@ class QsaPagedIndexScores(OpWrapper): **_TRUNK_POOL_NVFP4, ), ), - # ---- the target's call shapes: [*, 260] table, 2080 columns ------- Cell( why="decode: 1 row deep in history, [32, 260] table, visible-only", spec=_spec( @@ -155,59 +122,7 @@ class QsaPagedIndexScores(OpWrapper): ), ), 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="decode MTP7: 1 request x 8 rows across a page edge, visible-only", - spec=_spec( - layout="decode", - gen_positions=(6012,), - rows_per_request=8, - slots=32, - pages=_WIDTH, - only_visible_blocks=True, - context_rows=False, - branch="row_wise", - ), - ), - Cell( - why="decode MTP7: 32 slots x 8 rows = 256 rows, visible-only", + why="decode MTP7: 32 slots x 8 rows = 256 rows, the last decode-sized count", spec=_spec( layout="decode", gen_positions=tuple(100 + (499 * g) % 8000 for g in range(32)), @@ -219,67 +134,6 @@ class QsaPagedIndexScores(OpWrapper): branch="row_wise", ), ), - Cell( - why="decode MTP7: 32 slots x 8 rows = 256 rows, -inf padded", - spec=_spec( - layout="decode", - gen_positions=tuple(100 + (499 * g) % 8000 for g in range(32)), - rows_per_request=8, - slots=32, - pages=_WIDTH, - only_visible_blocks=False, - 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( @@ -293,31 +147,7 @@ class QsaPagedIndexScores(OpWrapper): ), ), 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", + why="mixed step: 195-token chunk + 31 MTP decodes = 257 rows: query-tiled", spec=_spec( layout="mixed", requests=32, @@ -329,32 +159,6 @@ class QsaPagedIndexScores(OpWrapper): 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( @@ -368,79 +172,6 @@ class QsaPagedIndexScores(OpWrapper): 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: 128 rows per query tile", - spec=_small(300, 5, heads=1, context_rows=True, branch="query_tile"), - ), - Cell( - why="300 rows x 2 heads: 64 rows per query tile", - spec=_small(300, 5, heads=2, context_rows=True, branch="query_tile"), - ), - Cell( - why="300 rows x 8 heads: 16 rows per query tile", - spec=_small(300, 5, heads=8, context_rows=True, branch="query_tile"), - ), - Cell( - why="300 rows x 16 heads: 8 rows per query tile", - spec=_small(300, 5, heads=16, context_rows=True, branch="query_tile"), - ), - Cell( - why="300 rows x 3 heads: does not divide 128, row-wise", - spec=_small(300, 5, heads=3, context_rows=True), - ), - Cell( - why="300 rows x 5 heads: does not divide 128, 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 = """ @@ -457,8 +188,8 @@ class QsaPagedIndexScores(OpWrapper): run: an uncompressed anchor scores as a raw index key. `rows == 0` returns an empty `[0, columns]` tensor without launching. - The op validates its own metadata loudly (10 `ValueError`s, pinned by the - test), so `is_valid` adds nothing. + The op validates its own metadata loudly (`ValueError`s; one per refusal + family is 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 @@ -481,10 +212,15 @@ class QsaPagedIndexScores(OpWrapper): visible block's, at a row whose position ends a group; that key is read again after the wait. - 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. + Every cell runs `only_visible_blocks=True` (the target never passes False), + 4 index heads x 128, compress ratio 4, 32-token pages, bf16 operands, int64 + positions and int32 request indices. + + Not certified: more than 8191 rows, tables wider than 260 pages, more than + 32 table rows, any other head geometry, compress ratio, page size or + operand dtype, and the `-inf` padded mode as a cell (the test gates it + against the reference outside the cells). Hard limits: not probed; the largest + certified extent is 8191 x 2080 output elements. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_store.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_store.py index a201309b62df..6bea19a8b576 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_store.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_store.py @@ -131,10 +131,7 @@ def _launch_rows(rows: int) -> Tuple[int, int]: request_dtype=torch.int32, position_dtype=torch.int64, key_layout="packed", - coords_layout="transposed", - meta_layout="contiguous", cache_layout="plain", - table_layout="dense", hole=None, ) @@ -146,7 +143,6 @@ def _launch_rows(rows: int) -> Tuple[int, int]: _WIDE = dict(slots=160, width=300) _DECODE_128 = tuple(((r * 72 + 9) % 9216, 1) for r in range(128)) _PREFILL_16K = ((0, 8192), (0, 8192)) -_MIXED_16K = ((0, 8192), (1024, 8064)) + tuple(((r * 71 + 5) % 9216, 1) for r in range(128)) #: 32 requests of one MTP7 verify block each (8 rows), at starts that put some blocks across a #: page edge and some inside one. @@ -192,56 +188,19 @@ class QsaPagedIndexStore(OpWrapper): CELLS: tuple[Cell, ...] = ( Cell( - why="decode: 8 rows at aligned, mid and page-crossing positions", - spec=_spec(_DECODE, pages=32), - ), - Cell( - why="decode MTP: 4 requests x 2 rows, a pair straddling a page", - spec=_spec(((10, 2), (31, 2), (63, 2), (5, 2)), pages=16), - ), - Cell( - why="decode MTP7: 1 request x 8 rows, a block straddling a page", - spec=_spec(((28, 8),), pages=8), - ), - Cell( - why="decode MTP7: 32 requests x 8 rows (256), blocks straddling pages", + why="decode MTP7: 32 requests x 8 rows (256), the last one-row-per-program " + "launch, blocks straddling pages", spec=_spec(_MTP7_32, pages=1200, slots=40), ), Cell( - why="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="decode: token_k a head of the fused projection row (110 x 128)", - spec=_spec(_DECODE, pages=32, key_layout="fused_row"), - ), - Cell( - why="prefill 1024 rows: token_k off the fused projection row", - spec=_spec(((0, 256),) * 4, pages=32, key_layout="fused_row"), - ), - 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"), + why="decode M=128 at positions up to 9153 (page column 286), token_k off the " + "fused projection row", + spec=_spec(_DECODE_128, pages=18432, key_layout="fused_row", **_WIDE), ), 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"), + why="prefill 16384 rows (max_num_tokens), 16 rows per program, token_k off the " + "fused projection row", + spec=_spec(_PREFILL_16K, pages=600, key_layout="fused_row", **_WIDE), ), Cell( why="NVFP4-KV side planes: index and position planes as strided views of one " @@ -255,81 +214,14 @@ class QsaPagedIndexStore(OpWrapper): position_strides=(24672, 3, 1), ), ), - 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="block table transposed: column stride != 1", - spec=_spec(_DECODE, pages=32, table_layout="transposed"), - ), - 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), - ), - # ---- the serving envelope: decode M=128, 16384-row steps ----------- - Cell( - why="decode M=128 at positions up to 9153 (page column 286)", - spec=_spec(_DECODE_128, pages=18432, **_WIDE), - ), - Cell( - why="decode M=128, token_k off the fused projection row", - spec=_spec(_DECODE_128, pages=18432, key_layout="fused_row", **_WIDE), - ), - Cell( - why="prefill 16384 rows: two 8192-token requests (max_num_tokens)", - spec=_spec(_PREFILL_16K, pages=600, **_WIDE), - ), - Cell( - why="prefill 16384 rows, token_k off the fused projection row", - spec=_spec(_PREFILL_16K, pages=600, key_layout="fused_row", **_WIDE), - ), - Cell( - why="mixed 16384 rows: two chunks, one mid-sequence, + 128 decode rows", - spec=_spec(_MIXED_16K, pages=18688, **_WIDE), - ), - Cell( - why="head_dim 64, 2 index-KV heads, 16-token pages", - spec=_spec(_DECODE, pages=48, hkv=2, head_dim=64, tokens_per_block=16), - ), - # ---- rows that resolve to no slot are skipped ------------------------ Cell( why="skip: table entry -1 (unallocated)", spec=_spec(_DECODE, pages=32, hole="unallocated_page"), ), - Cell( - why="skip: table entry past the pool", - spec=_spec(_DECODE, pages=32, hole="page_past_pool"), - ), - Cell( - why="skip: request index past the table", - spec=_spec(_DECODE, pages=32, hole="request_past_table"), - ), Cell( why="skip: negative request index", spec=_spec(_DECODE, pages=32, hole="negative_request"), ), - Cell( - why="skip: negative logical position", - spec=_spec(_DECODE, pages=32, hole="negative_position"), - ), - Cell( - why="skip: page column past the table", - spec=_spec(_DECODE, pages=32, hole="column_past_table"), - ), - Cell(why="zero rows: both planes unchanged", spec=_spec((), pages=8)), ) note = """ @@ -339,10 +231,10 @@ class QsaPagedIndexStore(OpWrapper): `get_index_k_buffer(i)` (`[pages, tpb, 1, 128]` bf16) and the shared `get_qsa_position_buffer()` (`[pages, tpb, 3]` int32, rewritten idempotently by every sparse layer). Every operand is read through its - own strides, all of them runtime arguments: `token_k` with any row stride - (the packed index projection's 640, the fused input projection's 14080, - contiguous), stride-0 or transposed coordinates, strided metadata, a - column-sliced or transposed table, a role view of the index plane. Under an + own strides, all of them runtime arguments; the cells drive the target's: + `token_k` off the packed index projection (row stride 640) and off the fused + input projection (14080), the `[T, 3]` transpose of the position buffer as + the coordinates, contiguous metadata and table. Under an NVFP4 KV pool the QSA cache manager packs each page's index-K roles and then its position page into one per-slot buffer, so both planes are strided views of one allocation: the trunk's index plane is `[pages, 32, 1, 128]` @@ -372,16 +264,22 @@ class QsaPagedIndexStore(OpWrapper): not specialized on, so a new batch or pool size does not recompile (strides still specialize on 1 and 16-divisibility). Up to 256 rows run one row per program (1 warp), above that 16 rows per program (4 warps). - CUDA-graph capturable (captured at 4 decode rows and at 256, 32 MTP7 verify - blocks); a replay follows positions rewritten in place. Repeating a call is + CUDA-graph capturable (captured at 256 rows, 32 MTP7 verify blocks); a + replay follows positions rewritten in place. Repeating a call is idempotent. Hard limits: every index is widened to int64 inside the kernel before any arithmetic; an int64 position at `2^32 + 5` lands past the table and is skipped rather than truncated. Rows are a grid axis: 2^31 programs. - Not certified: an fp8 index plane, multiple devices, concurrent streams - onto one pool, duplicate destinations. + Skips certified by a cell: an unallocated (`-1`) table entry and a negative + request index; the test also drives an int64 position past the table. The + other skip conditions share that mask and have no cell. + + Not certified: an fp8 index plane, keys other than bf16, any index geometry + other than 1 head x 128 at 32-token pages, strided metadata or tables, + stride-0 coordinates, multiple devices, concurrent streams onto one pool, + duplicate destinations. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store.py index e532a87f785f..9fb56a0e03c8 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store.py @@ -87,44 +87,18 @@ class QsaPagedKvStore(OpWrapper): ARCHS = frozenset({Arch.SM_103}) CELLS: tuple[Cell, ...] = ( + # The decode and prefill targets read k/v off the fused projection row; the MTP + # targets split them off a packed QKV row. Cell( - why="decode: 8 rows at aligned, mid and page-crossing positions", - spec=_spec(_DECODE, pages=64), - ), - Cell( - why="decode MTP: 4 requests x 2 rows, a pair straddling a page", - spec=_spec(((0, 2), (31, 2), (63, 2), (95, 2)), pages=32), - ), - Cell( - why="prefill: 1024 rows, 4 requests x 256 walking 8 pages each", - spec=_spec(((0, 256),) * 4, pages=160), - ), - Cell( - why="mixed: two context chunks and two decode rows, 66 rows", - spec=_spec(((0, 40), (0, 24), (200, 1), (7, 1)), pages=64), - ), - Cell( - why="fp8-e4m3 pool: the row rounded once", - spec=_spec(_DECODE, pages=64, cache_dtype=_F8), - ), - # The prefill target's call over an fp8 pool (kv_cache_dtype=fp8, unit scale). - Cell( - why="fp8 prefill: 1024 rows, 4 requests x 256 walking 8 pages each", - spec=_spec(((0, 256),) * 4, pages=160, cache_dtype=_F8), - ), - Cell( - why="fp8 mixed: two context chunks and two decode rows, 66 rows", - spec=_spec(((0, 40), (0, 24), (200, 1), (7, 1)), pages=64, cache_dtype=_F8), + why="decode: 8 rows off the fused projection row (row stride 14080), " + "at aligned, mid and page-crossing positions", + spec=_spec(_DECODE, pages=64, row_pad_heads=27), ), Cell( - why="fp8 decode MTP: 4 requests x 2 rows, a pair straddling a page", - spec=_spec(((0, 2), (31, 2), (63, 2), (95, 2)), pages=32, cache_dtype=_F8), + why="prefill 1024 rows off the fused projection row (row stride 14080)", + spec=_spec(((0, 256),) * 4, pages=160, row_pad_heads=27), ), # MTP7 verify blocks: 8 rows per request, so one block spans up to two pages. - Cell( - why="decode MTP7: 1 request x 8 rows, a block straddling a page", - spec=_spec(((28, 8),), pages=8), - ), Cell( why="decode MTP7: 32 requests x 8 rows (256), blocks straddling pages", spec=_spec(_MTP7_32, pages=1200, slots=40), @@ -137,77 +111,10 @@ class QsaPagedKvStore(OpWrapper): why="fp8 conversion: all 65536 bf16 bit patterns, as torch casts them", spec=_spec(((0, 128),), pages=8, cache_dtype=_F8, values="every_bf16"), ), - 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="decode: k/v as head slices of the fused projection row (row stride 14080)", - spec=_spec(_DECODE, pages=64, row_pad_heads=27), - ), - Cell( - why="prefill 1024 rows off the fused projection row (row stride 14080)", - spec=_spec(((0, 256),) * 4, pages=160, row_pad_heads=27), - ), - 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 = """ @@ -224,18 +131,19 @@ class QsaPagedKvStore(OpWrapper): - `request_indices` and `logical_positions` are contiguous: the kernel gets no stride for them. A strided view is read through its base pointer and rows land at other rows' slots, silently (`is_valid`). - Every other operand is read through its strides (packed QKV head slices, - head-major rows, pool role views, sliced or transposed tables). The - callable rejects loudly: mismatched or non-3-D/4-D K/V shapes, head - geometry, `k_cache.shape[2] != tokens_per_block`, row dtypes outside - bf16/fp16, pool dtypes outside bf16/fp16/fp8-e4m3, a non-int32 or - non-2-D table, non-integer or wrong-length metadata, CPU operands, an - empty pool or table. `rows == 0` returns before launching. + Every other operand is read through its strides; the cells drive k/v as + head slices of the packed QKV row and of the fused projection row, the + pool's K and V role views, and a contiguous table. The callable rejects + loudly: mismatched or non-3-D/4-D K/V shapes, head geometry, + `k_cache.shape[2] != tokens_per_block`, row dtypes outside bf16/fp16, + pool dtypes outside bf16/fp16/fp8-e4m3, a non-int32 or non-2-D table, + non-integer or wrong-length metadata, CPU operands, an empty pool or + table. `rows == 0` returns before launching. Strides and the pool/table extents are Triton constexprs: each distinct layout or pool size compiles its own kernel. CUDA-graph capturable - (captured at 8 decode rows and at 256, 32 MTP7 verify blocks); a replay - follows positions rewritten in place. + (captured at 256 rows, 32 MTP7 verify blocks); a replay follows positions + rewritten in place. Hard limits: both metadata values are cast to int32 inside the kernel (`.to(tl.int32)`); an int64 position at `2^32 + 5` was probed and lands @@ -243,7 +151,11 @@ class QsaPagedKvStore(OpWrapper): positions stay below 2^14. Not certified: duplicate destinations, multiple devices, a scaled fp8 - pool (there is no scale argument). + pool (there is no scale argument), fp16 rows, head geometry other than + this checkpoint's (D 256, Hkv 2, 32-token pages), int32 positions or + int64 request indices, head-major rows, sliced or transposed tables, + `rows == 0`, and of the rows that resolve to no slot any but an + unallocated (`-1`) table entry (the reference skips each). """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store_nvfp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store_nvfp4.py index 06b54e6a84d0..12a593a90d41 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store_nvfp4.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store_nvfp4.py @@ -380,76 +380,40 @@ class QsaPagedKvStoreNvfp4(OpWrapper): why="decode: 8 rows at aligned, mid-page and page-crossing positions", spec=_spec(_DECODE, pages=64), ), - Cell( - why="decode MTP draft 1: 4 requests x 2 rows, a pair straddling a page", - spec=_spec(((0, 2), (31, 2), (63, 2), (95, 2)), pages=32), - ), - Cell( - why="decode MTP draft 7: 4 requests x 8 rows across page edges", - spec=_spec(((0, 8), (28, 8), (60, 8), (1000, 8)), pages=64), - ), 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="every finite bf16 value, K in bit-pattern order and V shuffled across blocks", spec=_spec(((0, 128),), pages=8, values="finite_bf16"), ), Cell( - why="zeros, -0.0 and a 900.0 outlier inside blocks: zero scales and saturation", - spec=_spec(_DECODE, pages=64, values="edge"), - ), - Cell( - why="second-level scales g_k 0.37, g_v 2.5 (unit scale is what the target binds)", - spec=_spec(((0, 40), (0, 24), (200, 1), (7, 1)), pages=64, global_k=0.37, global_v=2.5), + why="zeros, -0.0 and a 900.0 outlier inside blocks, and a row on an unallocated " + "(-1) page: zero scales, saturation and the skip", + spec=_spec(_DECODE, pages=69, slots=8, values="edge", hole="unallocated_page"), ), Cell( why="the attention op's own append over the same rows (context, chunk, decode " "P=1 and P=3) at unit second-level scales: data and both scale layouts bitwise equal", spec=_spec(((0, 115), (0, 101), (0, 36)), pages=96, thop_append=True), ), - Cell( - why="the attention op's own append at second-level scales 0.37 and 2.5", - spec=_spec( - ((0, 115), (0, 101), (0, 36)), - pages=96, - thop_append=True, - global_k=0.37, - global_v=2.5, - ), - ), Cell( why="the engine's trunk layout: layer 11's planes among the 24 role pages of a slot", - spec=_spec(((0, 64), (1000, 2), (31, 2)), pages=96, slot_pages=24, layer=11), + spec=_spec(((0, 64), (1000, 2), (31, 2)), pages=96, slots=8, slot_pages=24, layer=11), ), # Triton compiles once per class of each plain-int argument -- equal to 1, a multiple # of 16, or anything else. The pool's slot count and the table's height are the two a - # served call varies; the cells above cover (multiple of 16, multiple of 16). + # served call varies; with the cells above (pool, height) is (16k, 16k), (other, 16k), + # (other, other) and (16k, other), and these two add table height 1. Cell( - why="specialization: slot count other, table height 32", - spec=_spec(_DECODE, pages=69), - ), - Cell( - why="specialization: slot count other, table height 1", + why="decode batch 1: slot count other, table height 1", spec=_spec(((0, 1),), pages=69, slots=1), ), Cell( - why="specialization: slot count other, table height 8", - spec=_spec(_DECODE, pages=69, slots=8), - ), - Cell( - why="specialization: slot count a multiple of 16, table height 1", + why="decode batch 1: slot count a multiple of 16, table height 1", spec=_spec(((0, 1),), pages=64, slots=1), ), - Cell( - why="specialization: slot count a multiple of 16, table height 8", - spec=_spec(_DECODE, pages=64, slots=8), - ), Cell( why="decode MTP graph: G = 1, 3 and 8 requests x 8 rows, replayed bitwise equal to eager", spec=_spec( @@ -458,33 +422,6 @@ class QsaPagedKvStoreNvfp4(OpWrapper): graph_batches=(1, 3, 8), ), ), - 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)), - 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 = """ @@ -502,15 +439,15 @@ class QsaPagedKvStoreNvfp4(OpWrapper): (page, slot) both store, unordered. - the dense path's reader assumes this layout: e2m1 nibbles HND per page, K scales linear, V scales in four-token groups. The attention op's - append, run over the same rows, leaves bitwise the same pool (cells), at - unit and non-unit second-level scales. + append, run over the same rows, leaves bitwise the same pool (a cell), + at the unit second-level scales the target binds. The callable rejects loudly (ValueError): non-matching or non-3-D k/v, a row dtype outside bf16/fp16, a head_dim that is not a power of two >= 32, a tokens_per_block that is not a multiple of 4, caches that are not uint8 of the stated shapes or whose last dim or scale head page is strided, a non-2-D or non-int32 table, metadata that is not one contiguous int32/int64 entry per row, a `kv_scale_orig_quant` that is not a contiguous fp32 `[3]`, operands on - another device. `rows == 0` returns before launching. + another device. `rows == 0` returns before launching (not certified). Every offset is formed in 64 bits: rows, heads, requests, positions and pages are widened before any product. CUDA-graph capturable; a replay follows rows, @@ -528,6 +465,10 @@ class (the k/v head slices of the packed QKV row, the planes' slot, head and Not finite values: a NaN or +-inf in a block is not certified (the e2m1 and e4m3 conversions saturate, the block maximum is the kernel's `max`). + + Not certified either: second-level scales other than 1.0, int32 positions + or int64 request indices, and of the rows that resolve to no slot any but + an unallocated (`-1`) table entry (the reference skips each). """ def __call__( @@ -577,7 +518,7 @@ def reference( Quantization in fp32 torch with IEEE reciprocals where the kernel uses `rcp.approx.ftz`; the two agree on every operand the cells drive (6.0, the - second-level scales, and every e4m3 scale value), which the cells' bitwise + unit second-level scale, and every e4m3 scale value), which the cells' bitwise gate is what pins. """ out = [ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa.py index be7d55e4991b..da6e70489838 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa.py @@ -116,125 +116,26 @@ class QsaPagedSparseGqa(OpWrapper): CELLS: tuple[Cell, ...] = ( # ---- generation: no prefill pair, split-K --------------------------- Cell( - why="decode: 1 row past the budget, split-K + merge, PDL", - spec=_decode(((4000, 1),)), - ), - Cell( - why="decode: 4 requests x 1 row, split-K + merge, PDL", - spec=_decode(((4000, 1), (4037, 1), (4074, 1), (4111, 1))), + why="decode MTP7: 32 slots x 8 rows = 256 rows, 512 programs, 4-way split-K", + spec=_decode(tuple((100 + (499 * g) % 8000, 8) for g in range(32))), ), Cell( why="decode: 1 row, gate off the fused row, in the merge, PDL", spec=_decode(((4000, 1),), q_layout="fused_row", gate="fused_row"), ), - Cell( - why="decode: 4 requests x 1 row, gate fused into the merge", - spec=_decode(tuple((4000 + 37 * i, 1) for i in range(4)), gate="fused_row"), - ), Cell( why="decode: 128 rows, 8-way split-K, gate fused into the merge", spec=_decode(((600, 1),) * 128, q_layout="fused_row", slots=160, gate="fused_row"), ), - Cell( - why="decode: 4 rows, a compact [rows, Hq, D] gate in the merge", - spec=_decode(((900, 1),) * 4, gate="heads"), - ), - 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="decode MTP7: 1 request x 8 rows, 32-way split-K", - spec=_decode(((4000, 8),)), - ), - Cell( - why="decode MTP7: 8 requests x 8 rows, 8-way split-K", - spec=_decode(tuple((4000 + 37 * i, 8) for i in range(8))), - ), - Cell( - why="decode MTP7: 32 slots x 8 rows = 256 rows, 512 programs, 4-way split-K", - spec=_decode(tuple((100 + (499 * g) % 8000, 8) for g in range(32))), - ), Cell( why="300 rows, no prefill pair: split-K, one split, no merge", spec=_decode(((600, 1),) * 300, launches=("splitk",), slots=384), ), Cell( - why="16-wide selection: one tile, the split-K kernel writes out", - spec=_decode(((600, 1),) * 4, launches=("splitk",), width=16), + why="holes: all-padding row, foreign request, unmapped page", + spec=_decode(((900, 1),) * 4, holes=True), ), # ---- 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), - ), - # The target's three call shapes over an fp8 pool (kv_cache_dtype=fp8, unit scale). - Cell( - why="fp8 decode MTP: 4 requests x 2 rows, split-K + merge", - spec=_decode(tuple((4000 + 37 * i, 2) for i in range(4)), **_FP8), - ), Cell( why="fp8 context 128 rows from 2032: pair passed, still split-K", spec=_decode(((2032, 128),), prefill_pair=True, **_FP8), @@ -243,34 +144,10 @@ class QsaPagedSparseGqa(OpWrapper): why="fp8 context 272 rows from 2032: causal-prefix prefill kernel", spec=_decode(((2032, 272),), launches=_PREFILL, prefill_pair=True, **_FP8), ), - Cell(why="fp8 mixed 257 rows: 1 chunk + 31 MTP decodes", spec=_mixed(32, (195,), **_FP8)), - Cell( - why="fp8 decode MTP7: 32 slots x 8 rows = 256 rows, 512 programs, 4-way split-K", - spec=_decode(tuple((100 + (499 * g) % 8000, 8) for g in range(32)), **_FP8), - ), - Cell( - why="int64 request indices", spec=_decode(((1000, 2),) * 4, request_dtype=torch.int64) - ), - Cell( - why="q as the head slice of the fused projection row (55 x 256), decode", - spec=_decode(((900, 1),) * 4, q_layout="fused_row"), - ), - Cell( - why="q off the fused projection row (55 x 256), 128 decode rows", - spec=_decode(((600, 1),) * 128, q_layout="fused_row", slots=160), - ), Cell( why="q off the fused projection row (55 x 256), 8191-row prefill", spec=_decode(((0, 8191),), launches=_PREFILL, prefill_pair=True, q_layout="fused_row"), ), - 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 = """ @@ -296,24 +173,26 @@ class docstring). QSA's producer -- `qsa_topk_select` with zero widens exactly to bf16 under a bf16 q), so it must hold K/V written at unit scale. - `D` must be a power of two (192 is refused); `Hq % Hkv == 0`. - The op validates its own inputs loudly (11 `ValueError`s pinned by the - test, plus a selection width the compress ratio cannot produce, refused - on the prefill kernel only), so `is_valid` adds nothing. `rows == 0` - returns an empty `[0, Hq, D]` tensor. + The op validates its own inputs loudly (`ValueError`s, one per guard + family 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. The decode site carries up to 8 rows per request (max_draft_len 7): the - cells run 1, 8 and 32 requests of 8 rows, which is 16, 128 and 512 - `rows * Hkv` programs, so the 32-way, 8-way and 4-way split-K launches. + cells run 32 requests of 8 rows, 512 `rows * Hkv` programs and a 4-way + split-K launch; the 1- and 128-row decode cells run more splits. 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. + `tokens_per_block` other than 32; more than 8191 rows; the split-K kernel + above 300 rows; an fp16 pool or q; standalone contiguous K/V buffers; + int64 request indices. + Hard limits: not probed; the largest certified extent is 8191 rows x 2051 + selected tokens through a [32, 260] table. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_decode.py index d53be3643322..b6180bcb7566 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_decode.py @@ -696,8 +696,6 @@ class QsaPagedSparseGqaDecode(OpWrapper): CELLS: tuple[Cell, ...] = ( Cell(why="decode, 1 row at position 8700, gated", spec=dict(pos=(8700,), gated=True)), - Cell(why="decode, 1 row at position 8700, no gate", spec=dict(pos=(8700,), gated=False)), - Cell(why="decode, 2 rows, gated", spec=dict(pos=(8700, 5000), gated=True)), Cell( why="decode, 4 rows, one a short sequence, gated", spec=dict(pos=(8700, 5000, 77, 3000), gated=True), @@ -726,6 +724,8 @@ class QsaPagedSparseGqaDecode(OpWrapper): 16-token tiles) and the merge sums the splits in split order (QSA's merge reduces them as a tree), so the fp32 values can differ in the last bits and the bf16 output by an ulp on some elements. + + Not certified: a call without `output_gate` (the target always passes one). """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4.py index 95fa751a5f84..1bfb5fd146b4 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4.py @@ -77,8 +77,8 @@ def _mixed(requests, context_lens, launches=_PREFILL, **overrides) -> dict: def _short(segments, slots, pool_pages) -> dict: - """Decode rows at low positions over a small pool: one cell per class of the - plain-int `num_pages` and `num_requests` Triton specializes on.""" + """Decode rows at low positions over a small pool whose page count and table height + pick the classes of the plain-int `num_pages` and `num_requests` Triton specializes on.""" return _decode(segments, slots=slots, pool_pages=pool_pages) @@ -869,68 +869,26 @@ class QsaPagedSparseGqaNvfp4(OpWrapper): why="decode: 1 row past the budget, 64-way 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 draft 3: 4 requests x 4 rows", - spec=_decode(tuple((4000 + 37 * i, 4) for i in range(4))), - ), - Cell( - why="decode: 15 rows (5 requests x 3), 32-way split-K", - spec=_decode(tuple((4000 + 37 * i, 3) for i in range(5))), - ), Cell( why="decode MTP draft 7: 16 requests x 8 rows (128 rows, 8-way)", spec=_decode(tuple((100 + (499 * g) % 8000, 8) for g in range(16))), ), - Cell( - why="decode MTP draft 7: 32 x 8 = 256 rows, 512 programs (4-way top)", - spec=_decode(tuple((100 + (499 * g) % 8000, 8) 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), - ), - Cell( - why="context 2 rows with the pair: 64-way split-K, PDL", - spec=_decode(((2032, 2),), prefill_pair=True), - ), 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: the 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 4096 rows from 0: rows see 1..4096 tokens", spec=_decode(((0, 4096),), launches=_PREFILL, prefill_pair=True), ), - Cell(why="mixed 257 rows: 1 chunk + 31 MTP decodes", spec=_mixed(32, (195,))), - Cell( - why="mixed 3806 rows: 9 chunks + 23 MTP decodes", - spec=_mixed(32, (1024, 900, 700, 500, 300, 200, 100, 30, 6)), - ), Cell( why="second-level scales g_k 0.37, g_v 2.5: the dequantize rounds to bf16", spec=_decode(tuple((4000 + 37 * i, 2) for i in range(4)), global_k=0.37, global_v=2.5), ), - Cell(why="fp16 q", spec=_decode(((900, 1),) * 4, q_dtype=torch.float16)), - Cell( - why="int64 request indices", spec=_decode(((1000, 2),) * 4, request_dtype=torch.int64) - ), - Cell(why="contiguous q", spec=_decode(((900, 1),) * 4, q_layout="contiguous")), Cell( why="holes: all-padding row, foreign request, unmapped page", spec=_decode(((900, 1),) * 4, holes=True), @@ -953,32 +911,18 @@ class QsaPagedSparseGqaNvfp4(OpWrapper): ), ), # Triton compiles once per class of each plain-int argument -- equal to 1, a multiple - # of 16, or anything else: the pool's page count and the table's height, and the - # split-K kernel's row count (rows 1, 15, 16, 128 and 256 above). + # of 16, or anything else: the pool's page count, the table's height (32 and 384 + # above) and the split-K kernel's row count (1, 16, 128 and 300 above). The two below + # add a pool of a multiple of 16 pages and table heights 1 and 8, so each class of + # each argument is reached, not every combination of them. Cell( why="specialization: pool of 100 pages (other), table height 1", spec=_short(((600, 1),), slots=1, pool_pages=100), ), - Cell( - why="specialization: pool of 100 pages (other), table height 8", - spec=_short(((600, 1),) * 4, slots=8, pool_pages=100), - ), - Cell( - why="specialization: pool of 100 pages (other), table height 16", - spec=_short(((600, 1),) * 4, slots=16, pool_pages=100), - ), - Cell( - why="specialization: pool of 96 pages (multiple of 16), table height 1", - spec=_short(((600, 1),), slots=1, pool_pages=96), - ), Cell( why="specialization: pool of 96 pages (multiple of 16), table height 8", spec=_short(((600, 1),) * 4, slots=8, pool_pages=96), ), - Cell( - why="specialization: pool of 96 pages (multiple of 16), table height 16", - spec=_short(((600, 1),) * 4, slots=16, pool_pages=96), - ), Cell( why="captured decode: G 1, 8 and 32 requests x 1 and 8 rows, replay == eager", spec=_decode((), slots=32, capture=dict(groups=(1, 8, 32), rows_per_request=(1, 8))), @@ -1016,7 +960,8 @@ class QsaPagedSparseGqaNvfp4(OpWrapper): Compiled variants: the kernels take the pool's page count (`num_pages`), the table's row count (`num_requests`) and the split-K kernel the step's row count as plain ints, and Triton compiles each apart at 1, at multiples of 16 - and otherwise. Every other plain int -- the strides, the table width 260 and + and otherwise; the cells reach each class of each of the three, not every + combination. Every other plain int -- the strides, the table width 260 and the selection's row stride 2051 -- is fixed at the served value the cells use. Preconditions: @@ -1034,7 +979,8 @@ class QsaPagedSparseGqaNvfp4(OpWrapper): Every row, request, page and position index is widened to 64 bits before any product. Not certified: head geometries other than 24/2 x 256; `tokens_per_block` other than 32; more than 8192 rows; capture of the - context/mixed site; a pool of one slot. + context/mixed site; a pool of one slot; an fp16 or unsliced q; int64 request + indices. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_prefill_compress.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_prefill_compress.py index 8efd444d874a..d70f62d41d8a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_prefill_compress.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_prefill_compress.py @@ -77,8 +77,7 @@ class QsaPrefillCompress(OpWrapper): Not idempotent: the anchor slot is the group's last input and its only output, so a second call over the same rows re-pools a compressed anchor - (measured: all 48 of 48 anchors moved, up to 2.94). Two rows naming one - anchor in a launch are harmless. + (measured: all 48 of 48 anchors moved, up to 2.94). Fusion boundary: pool, Gemma norm, position gather, rotation, store. The raw-key and position stores (`qsa_paged_index_store`, which must run @@ -97,63 +96,14 @@ class QsaPrefillCompress(OpWrapper): 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="decode MTP7: 1 request x 8 rows, two anchors, None", - spec=_spec(((4090, 8),), pages=160), - ), Cell( why="decode MTP7: 32 requests x 8 rows, two anchors each, None", spec=_spec(_MTP7_32, pages=1100, slots=40, rounding_flips=1), ), - Cell( - why="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"), - ), # Under an NVFP4 KV pool the position page shares each physical page with the index # pages, so both caches are strided views of one allocation: index page stride # 12 x 4096 + 192 = 49344 on the trunk (4288 on the draft), position strides @@ -166,11 +116,6 @@ class QsaPrefillCompress(OpWrapper): why="NVFP4 draft planes: one index role and the position page share each page", spec=_old(((0, 64),) * 2, pages=4, layout="shared_page", roles=1), ), - 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 = """ @@ -204,8 +149,8 @@ class QsaPrefillCompress(OpWrapper): `cos_sin`'s dtype is the rotation's precision: bf16 is accepted and was measured 1.4% of elements off by up to 1.56e-2 -- not certified. Text rows reduce mRoPE to plain RoPE bit for bit. Plain `@triton.jit`, no - autotune; CUDA-graph capturable at a fixed row count (captured at 128 rows - and at 256, 32 MTP7 verify blocks). + autotune; CUDA-graph capturable at a fixed row count (captured at 256 + rows, 32 MTP7 verify blocks). Under an NVFP4 KV pool both caches are strided views of one per-slot allocation (each physical page holds the index roles, then the position @@ -222,7 +167,11 @@ class QsaPrefillCompress(OpWrapper): [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. + or fp16 `cos_sin`, positions past the table's rows, multiple devices, + `tokens_per_block` other than 32 (no group straddles a page), eps other + than 1e-6, int32 positions or int64 request indices, a strided view of the + index cache, position plane or table other than the NVFP4 pool's shared + page, two rows naming one anchor in a launch, and `rows == 0`. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py index bc4d8345f915..1136b8dcee13 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py @@ -671,57 +671,22 @@ class QsaSparseGqaRowtile(OpWrapper): ARCHS = frozenset({Arch.SM_103}) CELLS: tuple[Cell, ...] = ( - Cell(why="context 300 rows from 0: rows see 1..300 tokens", spec=_spec([(0, 300)])), - Cell(why="context 272 rows from 2032: across the budget", spec=_spec([(2032, 272)])), Cell( - why="two requests (150 @1900, 150 @900): a group spans them", - spec=_spec([(1900, 150), (900, 150)]), + why="con 1 prefill: 8192 rows of one request from 0, across the budget", + spec=_spec([(0, 8192)]), ), - Cell(why="chunked continuation: 4096 rows from 4096", spec=_spec([(4096, 4096)])), Cell( why="sparse union: 4096 rows from 4096 ranking one score vector", spec=_spec([(4096, 4096)], shared_scores=True), ), - Cell(why="con 1 prefill: 8192 rows of one request from 0", spec=_spec([(0, 8192)])), - Cell(why="warm-up: 8191 rows of one request from 0", spec=_spec([(0, 8191)])), Cell( - why="max_num_tokens: 2 requests x 8192 rows from 0", - spec=_spec([(0, 8192), (0, 8192)]), - ), - Cell( - why="con 128 mixed: 3 chunks (16240 rows) + 52 decodes", + why="con 128 mixed: 3 chunks (16240 rows) + 52 decodes on the lone pass", spec=_spec([(0, 8192), (2192, 6000), (6144, 2048)] + _decodes(52), slots=128), ), - Cell(why="mixed: 1 chunk (195) + 31 MTP decodes", spec=_spec([(0, 195)] + _decodes(31, 2))), - Cell( - why="mixed: 4 chunks (8134 rows) + 28 MTP decodes", - spec=_spec([(0, 8000), (907, 100), (1814, 30), (2721, 4)] + _decodes(28, 2)), - ), - Cell( - why="mixed: 9 chunks (3760 rows) + 23 decodes", - spec=_spec( - [ - (0, 1024), - (907, 900), - (1814, 700), - (2721, 500), - (3628, 300), - (4535, 200), - (5442, 100), - (6349, 30), - (7256, 6), - ] - + _decodes(23) - ), - ), # fmt: skip Cell( why="int64 request indices and positions", spec=_spec([(1000, 300)], request_dtype=torch.int64), ), - Cell( - why="standalone contiguous K and V buffers, contiguous q", - spec=_spec([(900, 300)], cache_layout="contiguous", q_layout="contiguous"), - ), Cell( why="holes: all-padding row, foreign request, unmapped page", spec=_spec([(900, 40), (950, 40)], holes=True), diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_expand_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_expand_decode.py index 4d9bd57a18a3..470b7d0dab18 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_expand_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_expand_decode.py @@ -509,22 +509,10 @@ class QsaTopkExpandDecode(OpWrapper): why="decode, 1 row with ties at the threshold", spec=dict(rows=1, pos=(8700,), ties=True), ), - Cell( - why="decode, 2 rows, 2 MTP-style positions sharing a request", - spec=dict(rows=2, pos=(6001, 6002)), - ), Cell( why="decode, 1 row of scores within 5% of each other", spec=dict(rows=1, pos=(8700,), narrow=True), ), - Cell( - why="decode, 1 row of equal scores: block order decides", - spec=dict(rows=1, pos=(8700,), const=True), - ), - Cell( - why="decode, 1 row, 150 scores tied at the threshold", - spec=dict(rows=1, pos=(8700,), tied=150), - ), ) note = """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_select.py index a520d94dc9b9..295bfa6d32a2 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_select.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_select.py @@ -85,39 +85,6 @@ class QsaTopkSelect(OpWrapper): 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="decode MTP7: one request's 8 rows (6000..6007), NaN tail", - spec=dict( - rows=8, - cols=_SCORE_COLUMNS, - index_topk=_BLOCK_TOPK, - windows=("positions", tuple(range(6000, 6008))), - logits="relu", - tail="nan", - ), - ), Cell( why="decode MTP7: 32 requests x 8 rows = 256 rows, some below the budget, NaN tail", spec=dict( @@ -133,7 +100,7 @@ class QsaTopkSelect(OpWrapper): ), ), Cell( - why="context: 8191 rows at positions 0..8190 (warm-up capture)", + why="context: 8191 rows at positions 0..8190 (warm-up capture), the CUDA op", spec=dict( rows=8191, cols=_SCORE_COLUMNS, @@ -143,73 +110,6 @@ class QsaTopkSelect(OpWrapper): tail="nan", ), ), - Cell( - why="context: 1024 rows from position 4000 over 2048 cols, relu ties", - spec=dict( - rows=1024, - cols=2048, - index_topk=_BLOCK_TOPK, - windows=("chunk", 4000), - logits="relu", - tail=None, - ), - ), - Cell( - why="nonzero row_starts: offsets come back window-local", - spec=dict( - rows=4, - cols=512, - index_topk=32, - windows=("explicit", (0, 37, 100, 480), (512, 300, 140, 512)), - logits="randn", - tail=None, - ), - ), - Cell( - why="strided logits: column stride 2 takes the scalar load path", - spec=dict( - rows=5, - cols=512, - index_topk=32, - windows=("explicit", (0, 11, 0, 64, 7), (512, 400, 33, 512, 8)), - logits="randn", - tail=None, - strided=True, - ), - ), - Cell( - why="windows of 0, 1, k-1, k, 1 and full width: -1 suffix", - spec=dict( - rows=6, - cols=128, - index_topk=16, - windows=("explicit", (0, 0, 0, 5, 9, 0), (0, 1, 15, 21, 10, 128)), - logits="randn", - tail=None, - ), - ), - Cell( - why="budget wider than every window and than the matrix", - spec=dict( - rows=3, - cols=64, - index_topk=_BLOCK_TOPK, - windows=("explicit", (0, 0, 0), (64, 7, 0)), - logits="randn", - tail=None, - ), - ), - Cell( - why="row tier: 12287 rows, all on the insertion-sort variant", - spec=dict( - rows=6 * 2048 - 1, - cols=512, - index_topk=64, - windows=("random_ends",), - logits="randn", - tail=None, - ), - ), Cell( why="context: 16384 rows at positions 0..16383, NaN tail (CuTe DSL op)", spec=dict( @@ -221,17 +121,6 @@ class QsaTopkSelect(OpWrapper): tail="nan", ), ), - Cell( - why="row tier: 12416 rows, past 12287 the CuTe DSL op", - spec=dict( - rows=6 * 2048 + 128, - cols=512, - index_topk=64, - windows=("random_ends",), - logits="randn", - tail=None, - ), - ), ) note = """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_attn_decode_prologue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_attn_decode_prologue.py index 8338738fc361..854f7efb1d6c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_attn_decode_prologue.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_attn_decode_prologue.py @@ -356,9 +356,8 @@ class Qwen4ExpAttnDecodePrologue(OpWrapper): CELLS: tuple[Cell, ...] = ( Cell(why="decode, 1 row at position 8699 (a group's last token)", spec=dict(pos=(8699,))), - Cell(why="decode, 1 row at position 8700 (no compression)", spec=dict(pos=(8700,))), Cell( - why="decode, 4 rows, two group ends, one page boundary", + why="decode, 4 rows: three group ends (two on a page's last slot), one with no compression", spec=dict(pos=(8699, 8700, 4095, 31)), ), Cell( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_qk_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_qk_norm_rope.py index cf6f550b2e6c..8613a41e3c88 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_qk_norm_rope.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_qk_norm_rope.py @@ -976,33 +976,23 @@ class Qwen4ExpQkNormRope(OpWrapper): ARCHS = frozenset({Arch.SM_103}) CELLS: tuple[Cell, ...] = ( - Cell(why="decode M=1: one token, three distinct mRoPE rows", spec=dict(tokens=1, **CKPT)), - Cell(why="decode M=2", spec=dict(tokens=2, **CKPT)), - Cell(why="decode M=5: odd batch", spec=dict(tokens=5, **CKPT)), Cell( - why="decode M=16: last token count on the 1-head/2-warp launch", - spec=dict(tokens=16, **CKPT), + why="decode M=1: one token on the 1-head/2-warp launch, three distinct mRoPE rows", + spec=dict(tokens=1, **CKPT), ), Cell( why="decode M=17: first on the 2-head launch (13 main programs)", spec=dict(tokens=17, **CKPT), ), - Cell(why="decode M=33: one past a power of two", spec=dict(tokens=33, **CKPT)), - Cell(why="decode M=64: last on the 2-head launch", spec=dict(tokens=64, **CKPT)), Cell( why="decode M=65: first on the 8-head launch (last main program half masked)", spec=dict(tokens=65, **CKPT), ), # fmt: skip - Cell(why="decode M=127: odd, below max concurrency", spec=dict(tokens=127, **CKPT)), - Cell(why="decode M=128: max concurrency", spec=dict(tokens=128, **CKPT)), Cell( - why="decode M=128, positions up to the table end (262144 rows)", + why="decode M=128: max concurrency, positions up to the table end (262144 rows)", spec=dict(tokens=128, max_position=262144, **CKPT), ), - Cell(why="prefill 1023 tokens: last on the Triton kernel", spec=dict(tokens=1023, **CKPT)), Cell(why="prefill 1024 tokens: first on the CuTe kernel", spec=dict(tokens=1024, **CKPT)), - Cell(why="prefill 2048 tokens", spec=dict(tokens=2048, **CKPT)), - Cell(why="prefill 8193 tokens: one past the ISL", spec=dict(tokens=8193, **CKPT)), Cell(why="prefill 16384 tokens: max_num_tokens", spec=dict(tokens=16384, **CKPT)), ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/vision_rope_2d.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/vision_rope_2d.py index 7a4388ccbb9b..90917dac79e2 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/vision_rope_2d.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/vision_rope_2d.py @@ -30,24 +30,10 @@ MERGE = 2 THETA = 10000.0 -#: (frames, height, width) in patch units. The `serving_*` rows come from the -#: runtime-value dump of the mmmu gate (58.33): 170 segment sizes from 256 to -#: 12544 per frame, calls of up to 22144 rows. 62x90 is a grid this processor -#: emits; 112x112 the square form of the largest segment; 22144 the largest call. -GRIDS = { - "one_block": ((1, 2, 2),), - "tiny_image": ((1, 4, 4),), - "image": ((1, 32, 32),), - "prefill_image": ((1, 64, 64),), - "two_images": ((1, 16, 24), (1, 8, 12)), - "video": ((2, 4, 6),), - "video_and_image": ((3, 8, 8), (1, 4, 4)), - "serving_frame_5580": ((1, 62, 90),), - "serving_frame_12544": ((1, 112, 112),), - "serving_batch_22144": ((1, 112, 112), (1, 96, 100)), -} - -_DTYPES = (torch.bfloat16, torch.float16, torch.float32) +#: (frames, height, width) in patch units, from the runtime-value dump of the +#: mmmu gate (58.33): 170 segment sizes from 256 to 12544 per frame, calls of up +#: to 22144 rows. 62x90 is a grid this processor emits. +SERVING_FRAME = ((1, 62, 90),) def _rotate_half(x: torch.Tensor, half: int) -> torch.Tensor: @@ -100,73 +86,42 @@ class VisionRope2d(OpWrapper): # "packed_qkv" -- q and k as the `split` views of one [T, 3*heads, D] # projection, which is how the tower calls it. CELLS: tuple[Cell, ...] = ( - tuple( - Cell(why=f"{name} {str(dtype)[6:]}", spec=dict(grid=grid, dtype=dtype)) - for name, grid in GRIDS.items() - for dtype in _DTYPES - ) - + tuple( - Cell( - why=f"grid side {max(h, w)} ({h}x{w}) {str(dtype)[6:]}: largest angle", - spec=dict(grid=((1, h, w),), dtype=dtype), - ) - for h, w in ((2, 112), (2, 196), (2, 512), (2, 1024), (2, 6272), (6272, 2)) - for dtype in (torch.float32, torch.bfloat16) - ) - + tuple( - Cell( - why=f"head geometry {heads}x{head_dim}", - spec=dict(grid=((1, 8, 12),), dtype=torch.bfloat16, heads=heads, head_dim=head_dim), - ) - for heads, head_dim in ((8, 64), (4, 128), (2, 40)) - ) - + ( - Cell( - why="k carries its own head count (16 q / 2 k)", - spec=dict(grid=GRIDS["two_images"], dtype=torch.bfloat16, heads_k=2), - ), - Cell( - why="call site: q/k split views of packed qkv, 62x90", - spec=dict( - grid=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="call site: packed qkv views, video at the 32768-row ceiling", - spec=dict( - grid=((64, 16, 32),), - dtype=torch.bfloat16, - layout="packed_qkv", - grid_dtype=torch.int64, - ), - ), - Cell( - why="float32, video at the 32768-row ceiling, two frame groups of 128x128", - spec=dict(grid=((2, 128, 128),), dtype=torch.float32), - ), - 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="call site: q/k split views of packed qkv, 62x90", + spec=dict(grid=SERVING_FRAME, dtype=torch.bfloat16, layout="packed_qkv"), + ), + Cell( + why="call site: packed qkv views, video + image (two grid rows)", + spec=dict( + grid=((3, 8, 8), (1, 16, 24)), + dtype=torch.bfloat16, + layout="packed_qkv", + grid_dtype=torch.int64, ), - Cell( - why="theta 1e6: theta is an argument", - spec=dict(grid=GRIDS["tiny_image"], dtype=torch.float32, theta=1.0e6), + ), + Cell( + why="call site: packed qkv views, video at the 32768-row ceiling", + spec=dict( + grid=((64, 16, 32),), + dtype=torch.bfloat16, + layout="packed_qkv", + grid_dtype=torch.int64, ), + ), + Cell( + why="grid side 6272 (2x6272) float32: largest width angle, fp32 exposes the table", + spec=dict(grid=((1, 2, 6272),), dtype=torch.float32), + ), + Cell( + why="grid side 6272 (6272x2) bfloat16: largest height angle", + spec=dict(grid=((1, 6272, 2),), dtype=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)) ) note = """ @@ -177,9 +132,9 @@ class VisionRope2d(OpWrapper): `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 + Accuracy vs fp64 angles over the cells: fp32 <= 4.77e-7, 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`. @@ -189,8 +144,9 @@ class VisionRope2d(OpWrapper): 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 32768 rows. + Not certified: float16; mixed q/k dtypes or head counts; float8; + `spatial_merge_size` != 2; `theta` other than 1e4; grid sides past 6272 or + calls past 32768 rows. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/embedding/ngram_hash.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/embedding/ngram_hash.py index 21190a07f730..1892ee2c7cac 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/embedding/ngram_hash.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/embedding/ngram_hash.py @@ -67,31 +67,22 @@ class NgramHash(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec: batch, length, num_query_tokens; history is uniform random ids + # spec: batch, length, num_query_tokens; history is uniform random int64 ids # with the first `min(CONTEXT_LEN, length)` set to EOS when `lead_eos` # (default True), `eos_positions` forced to EOS, and `eos_density` of the - # remaining ids replaced by EOS; ids_dtype (int64). + # remaining ids replaced by EOS. + # + # Every cell is the call site's: one `[T, ngram_size]` window per packed + # token from `ngram_token_history`, num_query_tokens 1. CELLS: tuple[Cell, ...] = ( - _rows(1, 3, 1, "decode: 2 cached ids + 1 new, 1 row"), - _rows(8, 3, 1, "decode: 2 cached ids + 1 new, 8 rows"), _rows(4, 3, 1, "decode row whose cached context is EOS", eos_positions=(0, 1)), - _rows(1, 130, 128, "prefill chunk of 128"), - _rows(2, 130, 128, "prefill chunk of 128, EOS inside", eos_positions=(40, 41, 100)), - _rows(3, 1026, 1024, "prefill chunk of 1024, EOS runs", eos_positions=(7, 8, 9, 512)), - _rows(2, 64, 64, "whole history emitted", eos_positions=(16, 17)), - _rows(1, 2, 0, "empty query: [batch, 0, heads]"), - _rows(4, 6, 4, "4 query tokens per row (draft width)", eos_positions=(2,)), - # The call site: one `[T, ngram_size]` window per packed token from - # `ngram_token_history`, raw ids with EOS anywhere, num_query_tokens 1. _rows(1, 3, 1, "call site: 1 window, 1 token", lead_eos=False, eos_density=0.2), _rows( 32, 3, 1, "call site: 32 windows, decode / MTP verify", lead_eos=False, eos_density=0.2 ), - _rows(4096, 3, 1, "call site: 4096 windows, prefill", lead_eos=False, eos_density=0.1), _rows( 8192, 3, 1, "call site: 8192 windows, largest prefill", lead_eos=False, eos_density=0.05 ), - _rows(2, 34, 32, "int32 token_history", eos_positions=(10,), ids_dtype=torch.int32), ) note = """ @@ -103,11 +94,12 @@ class NgramHash(OpWrapper): Preconditions: `head_vocab_sizes`, `head_offsets` on `token_history`'s device (raises otherwise; a cpu `layer_multipliers` is accepted and - gives the identical result). `token_history` may be any integer dtype - (int32 and int64 identical). Negative ids are accepted and hashed as - ordinary ids, not padding. `num_query_tokens == 0` returns - `[batch, 0, heads]`. A single multiplier raises `ZeroDivisionError`. The - four silent mistakes are `is_valid`'s. + gives the identical result). Certified on an int64 `token_history`, the + window `ngram_token_history` returns. Negative ids are accepted and hashed + as ordinary ids, not padding. Longer histories and `num_query_tokens > 1` + (a chunk's tail) are exercised by the test beyond the cells, not certified. + A single multiplier raises `ZeroDivisionError`. The four silent mistakes + are `is_valid`'s. A wrong hash is a valid row: an extra shift, no EOS reset, an inclusive segment start, reversed multipliers, no offsets, swapped blocks all stay diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/bf16_gemv_pdl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/bf16_gemv_pdl.py index 82ab484714f2..a70d11d4f21e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/bf16_gemv_pdl.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/bf16_gemv_pdl.py @@ -322,21 +322,15 @@ class Bf16GemvPdl(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # Every shape is a qwen3.8-flash-next-nvfp4 tp1 decode projection. + # Every shape is a qwen3.8-flash-next-nvfp4 tp1 decode projection; the target calls + # it at one row (`_decode_projection`). CELLS: tuple[Cell, ...] = ( _cell("qwen3.8 decode M=1 router N=512", 1, 512, 2560), - _cell("qwen3.8 decode M=1 shared gate_up N=1280", 1, 1280, 2560), - _cell("qwen3.8 decode M=1 shared gate N=1", 1, 1, 2560), + _cell("qwen3.8 decode M=1 shared gate N=1, outputs per CTA halved to 1", 1, 1, 2560), _cell("qwen3.8 decode M=1 shared down K=640", 1, 2560, 640), - _cell("qwen3.8 decode M=3 shared down K=640, odd row count", 3, 2560, 640), - _cell("qwen3.8 decode M=1 GDN in-proj concat N=16480", 1, 16480, 2560), - _cell("qwen3.8 decode M=2 GDN in-proj concat", 2, 16480, 2560), - _cell("qwen3.8 decode M=1 GDN out / attention o, K=6144", 1, 2560, 6144), - _cell("qwen3.8 decode M=1 attention in-proj concat N=14080", 1, 14080, 2560), - _cell("qwen3.8 decode M=1 PLE key N=10240", 1, 10240, 2560), - _cell("qwen3.8 decode M=1 PLE value N=K=2560", 1, 2560, 2560), - _cell("qwen3.8 decode M=4 GDN out / attention o", 4, 2560, 6144), - _cell("qwen3.8 decode M=1 router, PDL off", 1, 512, 2560, pdl=False), + _cell("qwen3.8 decode M=1 GDN out / attention o, K=6144 shared-memory tier", 1, 2560, 6144), + _cell("qwen3.8 decode M=1 GDN in-proj concat N=16480, L2 prefetch grid", 1, 16480, 2560), + _cell("qwen3.8 decode M=4 (MAX_ROWS) GDN out / attention o", 4, 2560, 6144), ) note = """ @@ -358,9 +352,12 @@ class Bf16GemvPdl(OpWrapper): order differs from the fp32 reference. Hard limits - M rows per launch boundary MAX_ROWS=4 driven at 1, 2, 3, 4 (cells) + M rows per launch boundary MAX_ROWS=4 driven at 1 and 4 (cells) K a TILING key 2560, 6144, 640 all driven N any; the outputs per CTA halve until they divide N N=1 driven + + PDL off (`pdl=False`, or PDL disabled in the environment) is the same kernel + without the attribute; the test holds it bitwise to PDL on. """ def __call__(self, mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool = True) -> torch.Tensor: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/low_m_bf16_mm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/low_m_bf16_mm.py index 133d99940b7e..8405f6fb8ca4 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/low_m_bf16_mm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/low_m_bf16_mm.py @@ -144,42 +144,29 @@ class LowMBf16Mm(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # Every cell but the one off the table is a qwen3.8-flash-next-nvfp4 tp1 - # projection shape; `kind` pins the branch `plan` takes for it. + # Every cell is a qwen3.8-flash-next-nvfp4 tp1 decode projection at M >= 2 (one row + # takes `bf16_gemv_pdl` in the target); `kind` pins the branch `plan` takes for it. CELLS: tuple[Cell, ...] = ( - _cell("qwen3.8 decode M=1 GDN in_qkv: direct", 1, 10240, 2560, "direct"), - _cell("qwen3.8 decode M=2 GDN in_z: direct, 2-row tile", 2, 6144, 2560, "direct"), - _cell("qwen3.8 decode M=1 GDN in_a/in_b, N=48: direct", 1, 48, 2560, "direct"), + _cell( + "qwen3.8 decode M=2 GDN in_z: direct, 2-row tile, pdl on", + 2, + 6144, + 2560, + "direct", + pdl=True, + ), + _cell( + "qwen3.8 decode M=8 o_proj: split-K 3 over the cluster, pdl on", + 8, + 2560, + 6144, + "splitk", + pdl=True, + ), _cell("qwen3.8 decode M=32 GDN in_a/in_b: direct, 8-row tiles x4", 32, 48, 2560, "direct"), - _cell("qwen3.8 decode M=1 o_proj K=6144: direct, 256-thread K", 1, 2560, 6144, "direct"), - _cell("qwen3.8 decode M=8 o_proj: split-K 3 over the cluster", 8, 2560, 6144, "splitk"), _cell("qwen3.8 decode M=32 o_proj: split-K, mma_n 32", 32, 2560, 6144, "splitk"), - _cell("qwen3.8 decode M=3 attention q: split-K, M off-bucket", 3, 12288, 2560, "splitk"), - _cell("qwen3.8 decode M=8 attention q: split-K, mma_n 16", 8, 12288, 2560, "splitk"), - _cell("qwen3.8 decode M=16 router / k / v N=512: direct", 16, 512, 2560, "direct"), - _cell("qwen3.8 decode M=24 router / k / v: split-K 4", 24, 512, 2560, "splitk"), - _cell("qwen3.8 decode M=16 index proj N=640: split-K", 16, 640, 2560, "splitk"), - _cell("qwen3.8 decode M=4 shared gate_up: direct", 4, 1280, 2560, "direct"), - _cell("qwen3.8 decode M=8 shared gate_up + gate row: split-K", 8, 1296, 2560, "splitk"), - _cell("qwen3.8 decode M=1 shared gate N=1: direct", 1, 1, 2560, "direct"), - _cell("qwen3.8 decode M=1 GDN in-proj concat N=16480: direct", 1, 16480, 2560, "direct"), _cell("qwen3.8 decode M=2 GDN in-proj concat: cuBLAS branch", 2, 16480, 2560, "cublas"), - _cell( - "qwen3.8 decode M=2 attention in-proj concat N=14080: direct", 2, 14080, 2560, "direct" - ), - _cell("qwen3.8 decode M=3 attention in-proj concat: split-K", 3, 14080, 2560, "splitk"), - _cell("qwen3.8 decode M=8 attention in-proj concat: split-K", 8, 14080, 2560, "splitk"), - _cell("qwen3.8 decode M=16 attention in-proj concat: cuBLAS", 16, 14080, 2560, "cublas"), - _cell("qwen3.8 decode M=1 PLE value N=K=2560: direct", 1, 2560, 2560, "direct"), - _cell("qwen3.8 decode M=4 PLE value: direct, 32-thread K", 4, 2560, 2560, "direct"), - _cell("qwen3.8 decode M=8 PLE value: split-K 2", 8, 2560, 2560, "splitk"), - _cell("qwen3.8 decode M=1 shared down K=640: cuBLAS branch", 1, 2560, 640, "cublas"), - _cell("qwen3.8 decode M=16 attention q: cuBLAS branch", 16, 12288, 2560, "cublas"), _cell("qwen3.8 decode M=128: cuBLAS branch above M=32", 128, 10240, 2560, "cublas"), - _cell("qwen3.8 o_proj prefill-shaped M=8192: cuBLAS branch", 8192, 2560, 6144, "cublas"), - _cell("decode M=1, an (N, K) not in the table: cuBLAS", 1, 1024, 2048, "cublas"), - _cell("qwen3.8 pdl on, direct branch", 1, 10240, 2560, "direct", pdl=True), - _cell("qwen3.8 pdl on, split-K branch", 8, 2560, 6144, "splitk", pdl=True), ) note = """ @@ -203,7 +190,7 @@ class LowMBf16Mm(OpWrapper): Hard limits M FlashInfer branches only at M <= 32 (both kernels' policy) boundary 32 driven at 32 - (cells) and 128, 8192 (cuBLAS cells) + (cells) and 128 (cuBLAS cell) N*K largest FlashInfer-branch weight 16480*2560 = 4.2e7 elements boundary 2^31 not probed (48x below it) """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_gated_residual.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_gated_residual.py index fc016c9add86..38b3a8de5452 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_gated_residual.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_gated_residual.py @@ -68,41 +68,23 @@ class Qwen4ExpGatedResidual(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec: leading (token dims), hc_count, dtype (bf16), inject (True). Every - # cell uses synthetic weights with this checkpoint's measured element - # statistics, the normed stream reaching ~+-30 as the real one does. - 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), - ), + # spec: leading (token dims), inject (True). Every cell is this + # checkpoint's 4x2560 bf16 geometry with synthetic weights of its measured + # element statistics, the normed stream reaching ~+-30 as the real one does. + # Leading dims beyond one are `test_flattened_call_matches_batched`'s. + CELLS: tuple[Cell, ...] = tuple( + Cell(why=f"checkpoint 4x2560 bf16, {why}", spec=dict(leading=(t,))) + for t, why in ( + (1, "decode, 1 token"), + (8, "decode batch 8 / MTP draft rows"), + (512, "prefill 512 tokens"), + (8192, "prefill 8192 tokens, largest measured"), ) + ) + ( + Cell( + why="stack-final mixer: block_inject_weight None", + spec=dict(leading=(256,), inject=False), + ), ) note = """ @@ -133,9 +115,9 @@ class Qwen4ExpGatedResidual(OpWrapper): shared device: ~235 us at 1..512 tokens (overhead-bound), 475 us at 4096, 902 us at 8192. - Not certified: hidden other than 2560, lowrank other than 320, float32, - weight/activation dtype mixing. Hard limits: not probed (largest operand - 8192 x 10240 = 2^26.3 elements). + Not certified: hc_count other than 4, hidden other than 2560, lowrank + other than 320, float16 and float32, weight/activation dtype mixing. Hard + limits: not probed (largest operand 8192 x 10240 = 2^26.3 elements). """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first_cute.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first_cute.py index 6e813be14845..656d25581046 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first_cute.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first_cute.py @@ -424,18 +424,13 @@ class Qwen4ExpHcDecodeFirstCute(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec: rows (1), inject (True), moe (False: shared + gate_logits too), gate - # "column" [T, 1] | "flat" [T], fold (0: top-k of the finalize folded in, block_out - # then the permuted rows), drop (False: some slots -1), n_cols (324), n_rows (384: - # the padded weight), layout ("row_slice": hyper_input a row slice). + # spec: rows (1), inject (True), moe (False: shared + gate_logits [T, 1] too), fold + # (0: top-k of the finalize folded in, block_out then the permuted rows), drop + # (False: some slots -1), n_cols (324). The weight is the target's padded + # [width, 384] (columns past n_cols never read). CELLS: tuple[Cell, ...] = ( Cell(why="checkpoint 4x2560, 324 cols of 384, bf16, inject, 1 row", spec=dict(rows=1)), Cell(why="MoE combine, [T, 1] gate, 1 row", spec=dict(rows=1, moe=True)), - Cell(why="MoE combine, flat [T] gate, 1 row", spec=dict(rows=1, moe=True, gate="flat")), - Cell( - why="MoE combine into the stack-final mixer: 320 cols", - spec=dict(rows=1, moe=True, n_cols=320), - ), Cell(why="finalize fold, top-10, 1 row (con 1)", spec=dict(rows=1, moe=True, fold=10)), Cell( why="finalize fold into the stack-final mixer: 320 cols", @@ -446,9 +441,6 @@ class Qwen4ExpHcDecodeFirstCute(OpWrapper): spec=dict(rows=1, moe=True, fold=10, drop=True), ), Cell(why="norm-only site (stack start / after PLE)", spec=dict(rows=1, inject=False)), - Cell(why="stack-final mixer: 320 cols", spec=dict(rows=1, n_cols=320)), - Cell(why="hyper_input a row slice buf[::2]", spec=dict(rows=1, layout="row_slice")), - Cell(why="unpadded weight [width, 324]", spec=dict(rows=1, n_rows=324)), ) note = """ @@ -476,6 +468,9 @@ class Qwen4ExpHcDecodeFirstCute(OpWrapper): PDL: with `TRTLLM_ENABLE_PDL` (default on, sm_90+) the weight rows and the norm slice stream before `griddepcontrol.wait`, every activation read and every write after it, and dependents are released before the wait. + + Not certified: a flat `[T]` gate, a row-slice `hyper_input`, a weight + narrower than the padded 384 rows (accepted; no cell drives them). """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_tc_first.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_tc_first.py index 22416db4bd12..7de60f75d3e6 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_tc_first.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_tc_first.py @@ -561,28 +561,18 @@ class Qwen4ExpHcDecodeTcFirst(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec: rows, inject (True), hc_count (4), rank (320), hidden (2560), n_rows - # (384: the padded weight), dtype bf16 only. - CELLS: tuple[Cell, ...] = tuple( - Cell(why=f"checkpoint 4x2560, 324 cols of 384, inject, {why}", spec=dict(rows=t)) - for t, why in ( - (2, "decode 2 rows"), - (5, "decode 5 rows"), - (16, "decode 16 rows"), - (33, "decode 33 rows"), - (64, "decode 64 rows"), - (100, "decode 100 rows"), - (128, "decode 128 rows (max concurrency, MAX_ROWS)"), - ) - ) + ( - Cell(why="stack-final mixer: 320 cols, no inject, 2 rows", spec=dict(rows=2, inject=False)), + # spec: rows, inject (True). The checkpoint's geometry throughout: hc_count 4, + # rank 320, hidden 2560, the padded [384, K] weight, bf16. + CELLS: tuple[Cell, ...] = ( + Cell(why="checkpoint 4x2560, 324 cols of 384, inject, decode 2 rows", spec=dict(rows=2)), + Cell( + why="checkpoint 4x2560, 324 cols of 384, inject, decode 128 rows (MAX_ROWS)", + spec=dict(rows=128), + ), Cell( why="stack-final mixer: 320 cols, no inject, 128 rows", spec=dict(rows=128, inject=False), ), - Cell(why="unpadded weight [324, K], 8 rows", spec=dict(rows=8, n_rows=324)), - Cell(why="hc_count 2 at hidden 2560, 24 rows", spec=dict(rows=24, hc_count=2)), - Cell(why="1 row (below the dispatch range)", spec=dict(rows=1)), Cell(why="zero rows", spec=dict(rows=0)), ) @@ -614,6 +604,8 @@ class Qwen4ExpHcDecodeTcFirst(OpWrapper): against 8.6 us for cuBLAS split-K + reduce + lowrank_act; tile sweep BN 32 / 48 / 64 / 112 x KS 8 / 10 / 16, 48 x 16 fastest at every row count. + Not certified: hc_count other than 4, an unpadded weight (N = R + C), a + single row (below the dispatch range; accepted, no cell drives them). """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_inject_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_inject_norm.py index 6e61dbf47260..bfb9b7562991 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_inject_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_inject_norm.py @@ -780,73 +780,34 @@ class Qwen4ExpHcInjectNorm(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec: rows, inject (True), moe (False: shared + gate_logits too), gate - # "column" [T, 1] | "flat" [T], hc_count (4), hidden (2560), dtype (bf16), - # layout "contiguous" | "row_slice" (hyper_input = buf[::2]). + # spec: rows, inject (True), moe (False: shared + gate_logits [T, 1] too), fold + # (0: top-k of the finalize folded in), drop (False: some slots -1). The + # checkpoint's geometry throughout: hc_count 4, hidden 2560, bf16, contiguous rows. CELLS: tuple[Cell, ...] = ( - tuple( - Cell(why=f"checkpoint 4x2560 bf16, inject, {why}", spec=dict(rows=t)) - for t, why in ( - (1, "decode 1 token"), - (8, "decode batch 8"), - (33, "33 tokens, odd count"), - (128, "decode batch 128 (max concurrency)"), - (2048, "prefill chunk 2048"), - (8192, "prefill chunk 8192 (max_num_tokens)"), - ) - ) - + tuple( - Cell(why=f"norm only (no inject), {why}", spec=dict(rows=t, inject=False)) - for t, why in ((1, "decode 1 token"), (128, "decode 128"), (8192, "prefill 8192")) - ) - + tuple( - Cell(why=f"MoE combine, [T, 1] gate, {why}", spec=dict(rows=t, moe=True)) - for t, why in ( - (1, "decode 1 token"), - (8, "decode batch 8"), - (33, "33 tokens, odd count"), - (128, "decode batch 128 (max concurrency)"), - (2048, "prefill chunk 2048"), - (8192, "prefill chunk 8192"), - (16384, "prefill chunk 16384 (the ladder's max_num_tokens)"), - ) - ) - + tuple( - Cell( - why=f"MoE combine + finalize fold (top-10), {why}", - spec=dict(rows=t, moe=True, fold=10), - ) - for t, why in ( - (2, "decode batch 2"), - (33, "33 tokens, odd count"), - (128, "decode batch 128 (max concurrency)"), - ) - ) - + ( - Cell( - why="MoE combine + finalize fold, three slots mapped to -1, 8 rows", - spec=dict(rows=8, moe=True, fold=10, drop=True), - ), - Cell( - why="MoE combine, flat [T] gate, 17 rows", spec=dict(rows=17, moe=True, gate="flat") - ), - Cell( - why="MoE combine, float16 operands, 64 rows", - spec=dict(rows=64, moe=True, dtype=torch.float16), - ), - Cell(why="MoE combine, zero rows", spec=dict(rows=0, moe=True)), - ) - + ( - Cell( - why="hyper_input a row slice buf[::2], inject, 16 rows", - spec=dict(rows=16, layout="row_slice"), - ), - Cell(why="hc_count 2 at hidden 2560, 64 rows", spec=dict(rows=64, hc_count=2)), - Cell(why="hc_count 8 at hidden 2560, 64 rows", spec=dict(rows=64, hc_count=8)), - Cell(why="hidden 128 (not the checkpoint), 48 rows", spec=dict(rows=48, hidden=128)), - Cell(why="float16 operands, 512 rows", spec=dict(rows=512, dtype=torch.float16)), - Cell(why="zero rows", spec=dict(rows=0)), - ) + Cell( + why="checkpoint 4x2560 bf16, inject, decode batch 128 (eight-warp form)", + spec=dict(rows=128), + ), + Cell( + why="checkpoint 4x2560 bf16, inject, prefill chunk 8192 (two-warp form)", + spec=dict(rows=8192), + ), + Cell( + why="norm only (no inject: stack start / after PLE), decode 128", + spec=dict(rows=128, inject=False), + ), + Cell( + why="MoE combine, [T, 1] gate, prefill chunk 16384 (the ladder's max_num_tokens)", + spec=dict(rows=16384, moe=True), + ), + Cell( + why="MoE combine + finalize fold (top-10, CuTe kernel), decode batch 128", + spec=dict(rows=128, moe=True, fold=10), + ), + Cell( + why="MoE combine + finalize fold, three slots mapped to -1, 8 rows", + spec=dict(rows=8, moe=True, fold=10, drop=True), + ), ) note = """ @@ -900,6 +861,10 @@ class Qwen4ExpHcInjectNorm(OpWrapper): 16 rotating sites, us: 128 rows 7.00 -> 5.97, 64 rows 6.63 -> 5.25, 8 rows 3.59 -> 3.09; it releases dependents at its start as the Triton kernel does. + Not certified: hc_count other than 4, hidden other than 2560, float16 + operands, a flat `[T]` gate, row strides other than the width, zero rows + (the kernels take them; no cell drives them). + Hard limits: row offsets are int64 (`program_id(0)` widened); the grid is `rows * hc_count` programs along x (2^31 - 1); the column offset `stream * hidden + i` is int32, bounded by the width. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_lowrank_act.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_lowrank_act.py index 5c2b72a06ded..0dec54e742b6 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_lowrank_act.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_lowrank_act.py @@ -73,14 +73,10 @@ class Qwen4ExpHcLowrankAct(OpWrapper): for t, why in ( (1, "1 row"), (128, "128 rows"), - (2048, "prefill 2048"), (8192, "prefill 8192"), ) ) + ( Cell(why="stack-final mixer: no inject, 8192 rows", spec=dict(rows=8192, inject=False)), - Cell(why="unpadded logits [T, R + C], 33 rows", spec=dict(rows=33, logits_width=324)), - Cell(why="hc_count 8, 64 rows", spec=dict(rows=64, hc_count=8)), - Cell(why="float16 outputs, 64 rows", spec=dict(rows=64, dtype=torch.float16)), Cell(why="zero rows", spec=dict(rows=0)), ) @@ -89,6 +85,9 @@ class Qwen4ExpHcLowrankAct(OpWrapper): `PrefillTarget` (all 97), between the first projection and the up projection. Preconditions: `logits` fp32, unit column stride, at least `rank (+ hc_count)` columns; `hc_count` a power of two. Graph-capturable. + + Not certified: hc_count other than 4, rank other than 320, float16 + outputs. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean.py index 8c681b9b8572..60db120b3ac7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean.py @@ -835,19 +835,12 @@ class Qwen4ExpHcUpGateMean(OpWrapper): CELLS: tuple[Cell, ...] = tuple( Cell(why=f"checkpoint 4x2560, rank 320, {why}", spec=dict(rows=t)) for t, why in ( - (1, "decode 1 row"), - (32, "decode 32 rows"), - (128, "decode 128 rows"), - (2048, "prefill 2048"), - (8192, "prefill 8192"), + (1, "1 row, one partial tile"), + (128, "decode 128 rows, one full tile"), (16384, "prefill chunk 16384 (max_num_tokens)"), ) ) + ( Cell(why="ragged tail, 16320 + 70 rows (mixed step)", spec=dict(rows=16390)), - Cell(why="ragged tail, 1000 rows", spec=dict(rows=1000)), - Cell(why="hc_count 2, 300 rows", spec=dict(rows=300, hc_count=2)), - Cell(why="hc_count 8 (32-column tiles), 300 rows", spec=dict(rows=300, hc_count=8)), - Cell(why="hidden 128, rank 128, 200 rows", spec=dict(rows=200, hidden=128, rank=128)), Cell(why="zero rows", spec=dict(rows=0)), ) @@ -867,6 +860,9 @@ class Qwen4ExpHcUpGateMean(OpWrapper): `griddepcontrol.wait` and `low` / `normed` after it, so behind a producer that releases dependents early (`qwen4_exp_hc_decode_tc_first`) the weight read overlaps the producer. + + Not certified: hc_count other than 4, hidden other than 2560, rank other + than 320 (the kernel compiles for them; no cell drives them). """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean_fp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean_fp4.py index 601bfef7d350..49f2668d3374 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean_fp4.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean_fp4.py @@ -53,19 +53,11 @@ class Qwen4ExpHcUpGateMeanFp4(OpWrapper): Cell(why=f"checkpoint 4x2560, rank 320, {why}", spec=dict(rows=t)) for t, why in ( (2, "decode 2 rows (smallest fused-up batch)"), - (32, "decode 32 rows"), (128, "decode 128 rows (max_con)"), - (2048, "prefill 2048"), - (8192, "prefill 8192"), (16384, "prefill chunk 16384 (max_num_tokens)"), ) ) + ( - Cell(why="1 row (below the decode threshold)", spec=dict(rows=1)), Cell(why="ragged tail, 16320 + 70 rows (mixed step)", spec=dict(rows=16390)), - Cell(why="ragged tail, 1000 rows", spec=dict(rows=1000)), - Cell(why="hc_count 2, 300 rows", spec=dict(rows=300, hc_count=2)), - Cell(why="hc_count 8 (32-column tiles), 300 rows", spec=dict(rows=300, hc_count=8)), - Cell(why="hidden 128, rank 128, 200 rows", spec=dict(rows=200, hidden=128, rank=128)), Cell(why="zero rows", spec=dict(rows=0)), ) @@ -86,6 +78,9 @@ class Qwen4ExpHcUpGateMeanFp4(OpWrapper): instruction (`rcp.approx.ftz.f32` reciprocals, `cvt.rn.satfinite` to e4m3 and to e2m1, every multiply `mul.rn.f32`); the GPU test checks it bitwise against `torch.ops.trtllm.fp4_quantize` on the returned `mixed`. + + Not certified: hc_count other than 4, hidden other than 2560, rank other + than 320 (the kernel compiles for them; no cell drives them). """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_post_inject.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_post_inject.py index 679d4ce50521..aec5f3330ad4 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_post_inject.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_post_inject.py @@ -45,53 +45,25 @@ class Qwen4ExpPostInject(OpWrapper): # 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"), - ) + CELLS: tuple[Cell, ...] = tuple( + Cell( + why=f"checkpoint 4x2560 bf16, {why}", + spec=dict(leading=(t,), hc_count=CKPT_HC_COUNT, hidden=CKPT_HIDDEN), ) - + ( - Cell( - why="[2, 5, width] leading dims", - spec=dict(leading=(2, 5), hc_count=CKPT_HC_COUNT, hidden=CKPT_HIDDEN), - ), - Cell( - why="hyper_input a row slice buf[::2] of a wider buffer", - spec=dict( - leading=(8,), hc_count=CKPT_HC_COUNT, hidden=CKPT_HIDDEN, layout="row_slice" - ), - ), - Cell( - why="hyper_input a column slice buf[:, :width]", - spec=dict( - leading=(8,), hc_count=CKPT_HC_COUNT, hidden=CKPT_HIDDEN, layout="col_slice" - ), - ), - ) - + tuple( - Cell( - why=f"hc_count {c} read from inject, hidden 128", - spec=dict(leading=(48,), hc_count=c, hidden=128), - ) - for c in (1, 2, 3, 8) - ) - + tuple( - Cell( - why=f"{str(d).removeprefix('torch.')} operands, 64 tokens", - spec=dict(leading=(64,), hc_count=CKPT_HC_COUNT, hidden=CKPT_HIDDEN, dtype=d), - ) - for d in (torch.float16, torch.float32) + for t, why in ( + (1, "decode, 1 token"), + (64, "decode batch / MTP rows, 64 tokens"), + (8192, "prefill 8192 tokens, largest measured"), ) + ) + ( + 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"), + ), ) note = """ @@ -122,8 +94,9 @@ class Qwen4ExpPostInject(OpWrapper): broadcast operands take a non-vectorized path. An in-place form was not measured. - Not certified: float8 operands. Hard limits: not probed (the largest - measured operand is 8192 x 10240 = 2^26.3 elements). + Not certified: hc_count other than 4, hidden other than 2560, float16, + float32 and float8 operands, leading dims beyond one. Hard limits: not + probed (the largest measured operand is 8192 x 10240 = 2^26.3 elements). """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py index 8ab49bfa6aa6..8f993b017a57 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py @@ -137,79 +137,32 @@ class CuteDslNvfp4GatherGroupedGemmActFusion(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - CELLS: tuple[Cell, ...] = tuple( - cell - for tile in (128, 256) - for cell in ( - Cell( - why=f"tile {tile}: decode, 1 request without MTP", - spec=_ckpt(1, tile), - ), - Cell( - why=f"tile {tile}: 1 request at MTP verify width 2", - spec=_ckpt(2, tile), - ), - Cell( - why=f"tile {tile}: 1 request at MTP verify width 4", - spec=_ckpt(4, tile), - ), - Cell( - why=f"tile {tile}: decode, 8 requests", - spec=_ckpt(8, tile), - ), - Cell( - why=f"tile {tile}: decode at the 32-request ceiling", - spec=_ckpt(32, tile), - ), - Cell( - why=f"tile {tile}: 32 requests at MTP verify width 3", - spec=_ckpt(96, tile), - ), - Cell( - why=f"tile {tile}: 32 requests at MTP verify width 4; largest captured T", - spec=_ckpt(128, tile), - ), - Cell( - why=f"tile {tile}: 8192 rows, max_num_tokens prefill", - spec=_ckpt(8192, tile), - ), - Cell( - why=f"tile {tile}: 8192 rows all on 10 experts; {8192 // tile} full tiles each", - spec=_ckpt(8192, tile, dist="hot"), - ), - Cell( - why=f"tile {tile}: expert rows tile-1, tile, tile+1; the mn_limit edge", - spec=_ckpt(300, tile, dist="tile_edges"), - ), - Cell( - why=f"tile {tile}: 510 experts with one row each; every tile partial", - spec=_ckpt(51, tile, dist="round_robin"), - ), - Cell( - why=f"tile {tile}: MMA N 256 tactic, decode", - spec=_ckpt(1, tile, mma_n=256), - ), - Cell( - why=f"tile {tile}: MMA N 256 tactic, 128 rows", - spec=_ckpt(128, tile, mma_n=256), - ), - Cell( - why=f"tile {tile}: MMA N 256 tactic, 8192-row prefill", - spec=_ckpt(8192, tile, mma_n=256), - ), - Cell( - why=f"tile {tile}: global_sf 25x calibrated; block scales clip at 448", - spec=_ckpt(2048, tile, gsf="saturating"), - ), - Cell( - why=f"tile {tile}: tiny global_sf; subnormal and zero block scales", - spec=_ckpt(2048, tile, gsf="subnormal"), - ), - Cell( - why=f"tile {tile}: 7 all-zero token rows; zero block scales", - spec=_ckpt(64, tile, zero_tokens=7), - ), - ) + CELLS: tuple[Cell, ...] = ( + Cell( + why="tile 128: decode, 1 request without MTP", + spec=_ckpt(1, 128), + ), + Cell( + why="tile 128: MMA N 256 tactic, 32 requests at MTP verify width 4; largest captured T", + spec=_ckpt(128, 128, mma_n=256), + ), + Cell( + why="tile 128: 8192 rows, max_num_tokens prefill", + spec=_ckpt(8192, 128), + ), + Cell( + why="tile 128: expert rows tile-1, tile, tile+1; the mn_limit edge", + spec=_ckpt(300, 128, dist="tile_edges"), + ), + Cell( + why="tile 128: global_sf 25x calibrated; block scales clip at 448", + spec=_ckpt(2048, 128, gsf="saturating"), + ), + Cell( + why="tile 128: tiny global_sf and 7 all-zero token rows; subnormal and zero " + "block scales", + spec=_ckpt(2048, 128, gsf="subnormal", zero_tokens=7), + ), ) note = """ @@ -228,13 +181,15 @@ class CuteDslNvfp4GatherGroupedGemmActFusion(OpWrapper): Call sites: the engine's CuteDSL MoE (`run_moe_nvfp4_impl`) and the target's routed-expert step, eager for prefill and mixed batches (T up to max_num_tokens 8192) and captured for decode (T = requests x MTP verify width, - up to 32 x 4 = 128). Tile 128 is the target's; 256 is certified as well. + up to 32 x 4 = 128). Tile 128 is the target's and the cells' tile; 256 is + certified through the per-pair invariance below. Tactics. The op asks the AutoTuner for an MMA tile of (tile_size, 128) or (tile_size, 256) and with no tuning entry falls back to (tile_size, 128); serving warmup tunes per row bucket and may pick either. Every pair's codes and block scales are bitwise identical across tile 128 / 256 x MMA N 128 / 256 - x the fallback, so cells driven at one tactic certify the others. + x the fallback (the test checks it at 128 and 8192 rows), so cells driven at + one tactic certify the others. Determinism. A (token, expert) pair's output row depends on that pair alone: it is bitwise unchanged by reordering an expert's rows (moe_sort orders them by @@ -254,8 +209,10 @@ class CuteDslNvfp4GatherGroupedGemmActFusion(OpWrapper): 0. A block whose scale is 0 dequantizes to zero whatever its codes hold, and the test's compare does not read them. - Preconditions the caller owns (each wrong preparation ran without error, with - finite scales, and put most block scales and codes outside the test's gate): + Preconditions the caller owns (a wrong preparation of each operand -- the + weight's rows, its scales, the activation's scales, alpha -- runs without + error, with finite scales, and puts most block scales and codes outside the + test's gate): - `weight` / `weight_scale` in the layout the class docstring states: not trtllm-gen's row interleave, up rows first, scales swizzled after the same row interleave as the weight, alpha not shifted by an expert. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_grouped_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_grouped_gemm.py index c3398b0e579f..47359870c230 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_grouped_gemm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_grouped_gemm.py @@ -122,87 +122,33 @@ class CuteDslNvfp4GroupedGemm(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - CELLS: tuple[Cell, ...] = tuple( - cell - for tile in (128, 256) - for cell in ( - Cell( - why=f"tile {tile}: decode, 1 request without MTP", - spec=_ckpt(1, tile), - ), - Cell( - why=f"tile {tile}: 1 request at MTP verify width 2", - spec=_ckpt(2, tile), - ), - Cell( - why=f"tile {tile}: 1 request at MTP verify width 4", - spec=_ckpt(4, tile), - ), - Cell( - why=f"tile {tile}: decode, 8 requests", - spec=_ckpt(8, tile), - ), - Cell( - why=f"tile {tile}: decode at the 32-request ceiling", - spec=_ckpt(32, tile), - ), - Cell( - why=f"tile {tile}: 32 requests at MTP verify width 3", - spec=_ckpt(96, tile), - ), - Cell( - why=f"tile {tile}: 32 requests at MTP verify width 4; largest captured T", - spec=_ckpt(128, tile), - ), - Cell( - why=f"tile {tile}: 8192 rows, max_num_tokens prefill", - spec=_ckpt(8192, tile), - ), - Cell( - why=f"tile {tile}: 8192 rows all on 10 experts; {8192 // tile} full tiles each", - spec=_ckpt(8192, tile, dist="hot"), - ), - Cell( - why=f"tile {tile}: expert rows tile-1, tile, tile+1; the tile edge", - spec=_ckpt(300, tile, dist="tile_edges"), - ), - Cell( - why=f"tile {tile}: 510 experts with one row each; every tile mostly padding", - spec=_ckpt(51, tile, dist="round_robin"), - ), - Cell( - why=f"tile {tile}: MMA N 256 tactic, decode", - spec=_ckpt(1, tile, mma_n=256), - ), - Cell( - why=f"tile {tile}: MMA N 256 tactic, 128 rows", - spec=_ckpt(128, tile, mma_n=256), - ), - Cell( - why=f"tile {tile}: MMA N 256 tactic, 8192-row prefill", - spec=_ckpt(8192, tile, mma_n=256), - ), - Cell( - why=f"tile {tile}: cluster N 2 tactic, 128 rows", - spec=_ckpt(128, tile, cluster_n=2), - ), - Cell( - why=f"tile {tile}: MMA N 256 and cluster N 2 tactic, 8192-row prefill", - spec=_ckpt(8192, tile, mma_n=256, cluster_n=2), - ), - Cell( - why=f"tile {tile}: FC1 global_sf 25x calibrated; input scales clip at 448", - spec=_ckpt(2048, tile, gsf="saturating"), - ), - Cell( - why=f"tile {tile}: tiny FC1 global_sf; subnormal and zero input scales", - spec=_ckpt(2048, tile, gsf="subnormal"), - ), - Cell( - why=f"tile {tile}: 7 all-zero token rows; zero input rows", - spec=_ckpt(64, tile, zero_tokens=7), - ), - ) + CELLS: tuple[Cell, ...] = ( + Cell( + why="tile 128: decode, 1 request without MTP; the fallback tactic", + spec=_ckpt(1, 128), + ), + Cell( + why="tile 128: cluster N 2 tactic, 32 requests at MTP verify width 4; largest " + "captured T", + spec=_ckpt(128, 128, cluster_n=2), + ), + Cell( + why="tile 128: MMA N 256 and cluster N 2 tactic, 8192-row max_num_tokens prefill", + spec=_ckpt(8192, 128, mma_n=256, cluster_n=2), + ), + Cell( + why="tile 128: expert rows tile-1, tile, tile+1; the tile edge", + spec=_ckpt(300, 128, dist="tile_edges"), + ), + Cell( + why="tile 128: FC1 global_sf 25x calibrated; input scales clip at 448", + spec=_ckpt(2048, 128, gsf="saturating"), + ), + Cell( + why="tile 128: tiny FC1 global_sf and 7 all-zero token rows; subnormal and zero " + "input scales", + spec=_ckpt(2048, 128, gsf="subnormal", zero_tokens=7), + ), ) note = """ @@ -223,8 +169,8 @@ class CuteDslNvfp4GroupedGemm(OpWrapper): finalize is off) and the target's routed-expert step, both followed by `moe_unpermute_inplace`; eager for prefill and mixed batches (T up to max_num_tokens 8192) and captured for decode (T = requests x MTP verify - width, up to 32 x 4 = 128). Tile 128 is the target's; 256 is certified as - well. + width, up to 32 x 4 = 128). Tile 128 is the target's and the cells' tile; + 256 is certified through the per-pair invariance below. Numerics: each written element equals the reference's `out` bit for bit on every cell, on the real FC1 output at 8192 tokens, and on the checkpoint's own @@ -236,8 +182,8 @@ class CuteDslNvfp4GroupedGemm(OpWrapper): 128, 2), and with no tuning entry falls back to ((tile_size, 128), (tile_size // 128, 1)); serving warmup tunes per row bucket and may pick any of the four. Every pair's row is bitwise identical across tile 128 / 256 x - the four tactics x the fallback, so cells driven at one tactic certify the - others. + the four tactics x the fallback (the test checks it at 128 and 8192 rows), so + cells driven at one tactic certify the others. Determinism. A (token, expert) pair's output row depends on that pair's input row and its expert alone: it is bitwise unchanged by moving it to another row @@ -257,10 +203,10 @@ class CuteDslNvfp4GroupedGemm(OpWrapper): Preconditions the caller owns: - `weight` / `weight_scale` in the layout the class docstring states, and - `input_scale` swizzled. Each wrong preparation ran without error, with - finite output, and put most of the pair rows' elements outside the test's - gate: trtllm-gen's 32-row shuffle, weight scales left linear, input scales - linear, alpha shifted by one expert, one alpha for the whole layer. + `input_scale` swizzled. A wrong preparation of each runs without error, + with finite output, and puts most of the pair rows' elements outside the + test's gate: trtllm-gen's 32-row shuffle of the weight, input scales + linear, one alpha for the whole layer. - `alpha` per expert, `max(down_proj.input_scale) * weight_scale_2[e]`, with the maximum the value FC1's `global_sf` inverts. - `tile_idx_to_group_idx` / `num_non_exiting_tiles` from the same sort as diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_permuted.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_permuted.py index a925831d4631..9c381e0b0b65 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_permuted.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_permuted.py @@ -53,10 +53,12 @@ class Fp4BlockScaleMoeRenormPermuted(OpWrapper): ARCHS = frozenset({Arch.SM_103}) CELLS: tuple[Cell, ...] = ( - Cell(why="_moe_block decode, finalize folded into the next site: one token", spec=_ckpt(1)), - Cell(why="4 rows: the last count on the block routing kernel", spec=_ckpt(4)), - Cell(why="5 rows: past the block kernel's row limit", spec=_ckpt(5)), - Cell(why="16 rows: decode_first's MAX_ROWS", spec=_ckpt(16)), + Cell( + why="_moe_block decode, finalize folded into the next site: one token; the block " + "routing kernel", + spec=_ckpt(1), + ), + Cell(why="16 rows: decode_first's MAX_ROWS; the cluster routing kernel", spec=_ckpt(16)), Cell(why="128 rows: the con=128 decode batch", spec=_ckpt(128)), Cell(why="coarse 3-level logits: ties decide most slots", spec=_ckpt(64, dist="coarse")), ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_routed.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_routed.py index f74817c05d6a..d9d50991098f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_routed.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_routed.py @@ -62,15 +62,13 @@ class Fp4BlockScaleMoeRenormRouted(OpWrapper): ARCHS = frozenset({Arch.SM_103}) CELLS: tuple[Cell, ...] = ( - Cell(why="_moe_block decode: one token", spec=_ckpt(1)), - Cell(why="4 rows: the last count on the block routing kernel", spec=_ckpt(4)), - Cell(why="5 rows: past the block kernel's row limit", spec=_ckpt(5)), - Cell(why="8 rows: a decode CUDA-graph batch", spec=_ckpt(8)), - Cell(why="32 rows: a decode CUDA-graph batch", spec=_ckpt(32)), - Cell(why="128 rows: the con=128 decode batch", spec=_ckpt(128)), - Cell(why="1024 rows: a prefill chunk", spec=_ckpt(1024)), + Cell(why="_moe_block decode: one token; the block routing kernel", spec=_ckpt(1)), + Cell(why="128 rows: the con=128 decode batch; the cluster routing kernel", spec=_ckpt(128)), Cell(why="8192 rows: max_num_tokens of the identity config", spec=_ckpt(8192)), - Cell(why="16384 rows: max_num_tokens of the tuned config", spec=_ckpt(16384)), + Cell( + why="16384 rows: max_num_tokens of the tuned config; past the 8192 tuning bucket", + spec=_ckpt(16384), + ), Cell(why="coarse 3-level logits: ties decide most slots", spec=_ckpt(64, dist="coarse")), Cell( why="fp32 logits: the runner's other accepted logits dtype", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_moe_renorm_cute_dsl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_moe_renorm_cute_dsl.py index a02b4e47257b..85cbcb72ae9f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_moe_renorm_cute_dsl.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_moe_renorm_cute_dsl.py @@ -48,14 +48,18 @@ class Fp4MoeRenormCuteDsl(OpWrapper): ARCHS = frozenset({Arch.SM_103}) CELLS: tuple[Cell, ...] = ( - Cell(why="1024 rows: a prefill chunk", spec=_ckpt(1024)), - Cell(why="4096 rows: where the target switches to this entry", spec=_ckpt(4096)), - Cell(why="8192 rows", spec=_ckpt(8192)), + Cell( + why="4096 rows: where the target switches to this entry; the sort's cluster kernel", + spec=_ckpt(4096), + ), Cell( why="16377 rows: a con=128 mixed step (16256 context + 121 generation)", spec=_ckpt(16377), ), - Cell(why="16384 rows: max_num_tokens of the tuned config", spec=_ckpt(16384)), + Cell( + why="16384 rows: max_num_tokens of the tuned config; the sort's cooperative kernel", + spec=_ckpt(16384), + ), ) note = """ @@ -81,7 +85,8 @@ class Fp4MoeRenormCuteDsl(OpWrapper): (prefill steps run eagerly). Hard limits - T certified 1024..16384 + T certified 4096..16384 (the target calls it from + `_CUTE_DSL_MOE_MIN_ROWS` = 4096 rows) num_experts 512, top_k 10 (the routed entry's CERTIFIED_PAIRS) """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp8_block_scale_moe_runner.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp8_block_scale_moe_runner.py index 9fe9ee74e6a4..09cbecdc20f0 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp8_block_scale_moe_runner.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp8_block_scale_moe_runner.py @@ -101,83 +101,29 @@ class Fp8BlockScaleMoeRunner(OpWrapper): ARCHS = frozenset({Arch.SM_103}) CELLS: tuple[Cell, ...] = ( - # The token column at the checkpoint geometry. The cold-cache fallback - # tactic is keyed on the token count, and rows per active expert - # (T * 10 / 512) is what moves the kernel's tile choice. + # The token column at the checkpoint geometry, the only one the draft layer + # calls. The cold-cache fallback tactic is keyed on the token count, and rows + # per active expert (T * 10 / 512) is what moves the kernel's tile choice. Cell( why="MTPLayer._moe decode: one token, 10 of 512 experts", spec=dict(tokens=1, **_CKPT) ), - Cell(why="2 tokens: up to 20 experts, one row each", spec=dict(tokens=2, **_CKPT)), - Cell(why="8 tokens: ~0.16 rows per expert", spec=dict(tokens=8, **_CKPT)), - Cell(why="16 tokens: ~0.3 rows per expert", spec=dict(tokens=16, **_CKPT)), Cell( why="20 tokens == H/128: the activation scale buffer is square", spec=dict(tokens=20, **_CKPT), ), Cell( - why="64 tokens: ~1.25 rows per expert, most experts active", - spec=dict(tokens=64, **_CKPT), + why="256 tokens: ~5 rows per expert, past the small-T routing kernel", + spec=dict(tokens=256, **_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 of the identity config", + why="8192 tokens: max_num_tokens of the identity config; the last tuning bucket", spec=dict(tokens=8192, **_CKPT), ), Cell( - why="16384 rows: max_num_tokens of the tuned config", + why="16384 rows: max_num_tokens of the tuned config; past the last tuning bucket", spec=dict(tokens=16384, **_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 = """ @@ -198,9 +144,10 @@ class Fp8BlockScaleMoeRunner(OpWrapper): twice at 256 (the token count picks the routing kernel). The window: local slot i answers for global id `local_expert_offset + i`. - Nothing checks the window against `num_experts`; ids outside it (negative, - `>= num_experts`) are dropped, a token with no slot in it comes back - bitwise zero, and a window starting at `num_experts` returns all zeros. + Nothing checks the window against `num_experts`; a token with no slot in it + comes back bitwise zero, and a window starting at `num_experts` returns all + zeros. Ids outside the routing space (negative, `>= num_experts`) were + measured to be dropped; the router never emits one, and no cell certifies it. Four 128-wide windows of 512 summed match the full call inside the test's band but not bitwise (each rounds its partial to bf16). @@ -226,11 +173,12 @@ class Fp8BlockScaleMoeRunner(OpWrapper): `tune_max_num_tokens` (default 8192, which the target leaves): past it a call runs the 8192 bucket's tactic; a larger value makes the op append a bucket at that size - top_k <= 32 (raises past it); certified 1, 2, 6, 8, 10, 16 - E, H, I certified up to 512, 7168, 2048 + top_k <= 32 (raises past it); certified 10, the checkpoint's + E, H, I certified at the checkpoint's 512, 2560, 640 Not certified: the in-kernel routing path (`routing_logits` given), - multi-rank execution (the windows ran in one process), sm_100. + multi-rank execution (the windows ran in one process), geometries other + than the checkpoint's, sm_100. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fused_sigmoid_gate_mul_add.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fused_sigmoid_gate_mul_add.py index f05a5c6e24f1..473c7d9d2475 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fused_sigmoid_gate_mul_add.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fused_sigmoid_gate_mul_add.py @@ -59,30 +59,13 @@ class FusedSigmoidGateMulAdd(OpWrapper): why="_moe_block / MTPLayer._moe decode: one token, in place; 8 warps", spec=_ckpt(1), ), - Cell(why="2 tokens: the smallest multi-row grid", spec=_ckpt(2)), - Cell(why="17 tokens: an odd count", spec=_ckpt(17)), + Cell(why="17 tokens: an odd multi-row count", spec=_ckpt(17)), Cell(why="2048 tokens: prefill-shaped", spec=_ckpt(2048)), - Cell(why="fp16 activations", spec=_ckpt(33, dtype=torch.float16, gate_dtype=torch.float16)), - Cell(why="fp32 activations", spec=_ckpt(33, dtype=torch.float32, gate_dtype=torch.float32)), - Cell(why="fp32 gate against bf16 activations", spec=_ckpt(9, gate_dtype=torch.float32)), - Cell(why="gate as a flat [T] vector", spec=_ckpt(33, gate_layout="flat")), - Cell(why="gate as a [1, T] row", spec=_ckpt(33, gate_layout="row")), Cell( - why="gate as a strided column of a wider tensor", spec=_ckpt(33, gate_layout="strided") + why="gate as a strided column of a wider tensor, as the trunk's fused shared-expert " + "GEMM leaves it", + spec=_ckpt(33, gate_layout="strided"), ), - Cell(why="into an output buffer: both inputs untouched", spec=_ckpt(17, dest="output")), - Cell( - why="row-padded windows: every tensor's row stride exceeds hidden", - spec=_ckpt(33, dest="padded"), - ), - # Launch regimes: the column block is next_pow2(hidden) capped at 8192, - # and the warp count steps with it. - Cell(why="hidden 1: the 16-wide minimum block, 2 warps", spec=_ckpt(8, hidden=1)), - Cell(why="hidden 7: an odd width inside the minimum block", spec=_ckpt(8, hidden=7)), - Cell(why="hidden 640: a 1024 block, 4 warps", spec=_ckpt(8, hidden=640)), - Cell(why="hidden 6145: one 8192 block, mostly masked", spec=_ckpt(8, hidden=6145)), - Cell(why="hidden 12288: two column blocks", spec=_ckpt(8, hidden=12288)), - Cell(why="hidden 20480: three column blocks, masked tail", spec=_ckpt(8, hidden=20480)), ) note = """ @@ -98,7 +81,8 @@ class FusedSigmoidGateMulAdd(OpWrapper): 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. + Numerics: the default `assert_close` band holds for bf16, the dtype the target + calls it with; fp16 and fp32 run the same kernel and are not certified. bf16 at hidden 2560 is bitwise equal at 1-2 tokens and >= 99.993% bitwise at 17-8192 tokens, one ulp otherwise -- `tl.sigmoid`'s fast exponential. Wrong variants land 29,000-152,000x outside the bf16 band: no sigmoid, @@ -108,8 +92,10 @@ class FusedSigmoidGateMulAdd(OpWrapper): Hard limits num_tokens row offset widened to int64 in the kernel; largest run 8192 - Not certified: CUDA-graph capture, torch.compile, operands on two devices, - float8 activations. + Not certified: hidden sizes other than 2560 (the column block and warp count + step with it), an `output` buffer other than `shared_expert_output`, row-padded + operands, CUDA-graph capture, torch.compile, operands on two devices, float8 + activations. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_sort.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_sort.py index 8aa56756273e..ecfece5bfb9c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_sort.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_sort.py @@ -71,59 +71,33 @@ class MoeSort(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - CELLS: tuple[Cell, ...] = tuple( - cell - for tile in (128, 256) - for cell in ( - Cell( - why=f"tile {tile}: decode, 1 request without MTP; block kernel", - spec=_ckpt(1, tile), - ), - Cell( - why=f"tile {tile}: 4 rows, 1 request at MTP verify width 4; last block-kernel T", - spec=_ckpt(4, tile), - ), - Cell( - why=f"tile {tile}: 5 rows; first single-cluster-kernel T", - spec=_ckpt(5, tile), - ), - Cell( - why=f"tile {tile}: 32 rows, decode at the 32-request ceiling", - spec=_ckpt(32, tile), - ), - Cell( - why=f"tile {tile}: 128 rows, 32 requests at MTP verify width 4; largest captured T", - spec=_ckpt(128, tile), - ), - Cell( - why=f"tile {tile}: 8192 rows, max_num_tokens prefill; last single-cluster T", - spec=_ckpt(8192, tile), - ), - Cell( - why=f"tile {tile}: 8192 rows all to the same 10 experts; full multi-tile runs", - spec=_ckpt(8192, tile, "hot"), - ), - Cell( - why=f"tile {tile}: expert counts tile-1, tile, tile+1; the mn_limit clamp", - spec=_ckpt(300, tile, "tile_edges"), - ), - Cell( - why=f"tile {tile}: 510 experts with one row each; n_tiles == max_tiles", - spec=_ckpt(51, tile, "round_robin"), - ), - Cell( - why=f"tile {tile}: 8193 rows, past the envelope; first cooperative-kernel T", - spec=_ckpt(8193, tile), - ), - Cell( - why=f"tile {tile}: 471859 rows; last cooperative-kernel T on 152 SMs", - spec=_ckpt(471859, tile, "strided"), - ), - Cell( - why=f"tile {tile}: 471860 rows; first histogram+offsets T on 152 SMs", - spec=_ckpt(471860, tile, "strided"), - ), - ) + CELLS: tuple[Cell, ...] = ( + Cell( + why="tile 128: decode, 1 request without MTP; block kernel", + spec=_ckpt(1, 128), + ), + Cell( + why="tile 128: 128 rows, 32 requests at MTP verify width 4; largest captured T, " + "cluster kernel", + spec=_ckpt(128, 128), + ), + Cell( + why="tile 128: expert counts tile-1, tile, tile+1; the mn_limit clamp", + spec=_ckpt(300, 128, "tile_edges"), + ), + Cell( + why="tile 128: 8192 rows all to the same 10 experts; full multi-tile runs at the " + "last single-cluster T", + spec=_ckpt(8192, 128, "hot"), + ), + Cell( + why="tile 128: 8193 rows; first cooperative-kernel T", + spec=_ckpt(8193, 128), + ), + Cell( + why="tile 256: expert counts tile-1, tile, tile+1; the other tile sm_103 admits", + spec=_ckpt(300, 256, "tile_edges"), + ), ) note = """ @@ -181,8 +155,10 @@ class MoeSort(OpWrapper): <= 0 (the host divides by it). Certified only at the checkpoint's values: num_experts 512, top_k 10, all - experts local (offset 0), tile 128 and 256, fp32 scales. Expert parallelism, - other tiles and other expert counts are uncertified. + experts local (offset 0), tile 128 (the target's) and 256, fp32 scales, and + T up to 8193 (block, cluster and the first cooperative T). Expert + parallelism, other tiles, other expert counts and the histogram path are + uncertified. Hard limits (sm_103) T * top_k boundary 2^31 pairs (int32 product) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_unpermute_inplace.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_unpermute_inplace.py index 24ceefb36d4f..3fddd53c8ded 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_unpermute_inplace.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_unpermute_inplace.py @@ -106,35 +106,19 @@ class MoeUnpermuteInplace(OpWrapper): why="decode, 1 request without MTP; rows from the sort's block kernel", spec=_ckpt(1), ), - Cell(why="1 request at MTP verify width 2", spec=_ckpt(2)), Cell( - why="1 request at MTP verify width 4; last T the sort places in token order", - spec=_ckpt(4), + why="32 requests at MTP verify width 4; largest captured T, rows placed by the " + "sort's atomics", + spec=_ckpt(128), ), - Cell(why="decode, 8 requests; rows placed by the sort's atomics", spec=_ckpt(8)), - Cell(why="decode at the 32-request ceiling", spec=_ckpt(32)), - Cell(why="32 requests at MTP verify width 3", spec=_ckpt(96)), - Cell(why="32 requests at MTP verify width 4; largest captured T", spec=_ckpt(128)), Cell(why="8192 rows, max_num_tokens prefill", spec=_ckpt(8192)), - Cell( - why="8192 rows all on 10 experts; a token's rows in 10 runs of 8192", - spec=_ckpt(8192, dist="hot"), - ), - Cell( - why="510 experts with one row each; most permuted rows are padding", - spec=_ckpt(51, dist="round_robin"), - ), Cell( why="FC2 rows spanning 2^-30..2^30; terms 2^60 apart in one sum", spec=_ckpt(8192, values="wide"), ), Cell( - why="tile 256 sort layout, decode, 1 request", - spec=_ckpt(1, tile_tokens_dim=256), - ), - Cell( - why="tile 256 sort layout, 8192-row prefill", - spec=_ckpt(8192, tile_tokens_dim=256), + why="510 experts with one row each; most permuted rows are padding", + spec=_ckpt(51, dist="round_robin"), ), Cell( why="slots marked -1 (another rank's expert) skipped; a row of all -1 is 0", @@ -200,7 +184,7 @@ class MoeUnpermuteInplace(OpWrapper): rows; a CPU `output`. Certified only at the checkpoint's values: top_k 10, hidden 2560, bf16 rows and - output, fp32 scales, the sort's layouts at tile 128 and 256, every slot local + output, fp32 scales, the sort's layout at tile 128, every slot local except in the -1 cell. Hard limits (sm_103) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/renorm_moe_routing_op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/renorm_moe_routing_op.py index 261f6660ecbf..4944323b2283 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/renorm_moe_routing_op.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/renorm_moe_routing_op.py @@ -49,47 +49,19 @@ class RenormMoeRoutingOp(OpWrapper): 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="16384 rows: max_num_tokens of the tuned config", spec=_ckpt(16384)), Cell( - why="the default fp32 output dtype at the checkpoint router", + why="fp32 weights at the checkpoint router, as the CuteDSL experts take them", 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 = """ @@ -120,7 +92,8 @@ class RenormMoeRoutingOp(OpWrapper): 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. + hold the target's router (512 / top-10, bf16 logits, bf16 and fp32 weights) + and the test's sweep re-drives the boundaries of the rest of that space. Hard limits num_experts <= 512 (raises past it) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/flashinfer_gemma_rmsnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/flashinfer_gemma_rmsnorm.py index 57e25f06609e..4109a96da06f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/flashinfer_gemma_rmsnorm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/flashinfer_gemma_rmsnorm.py @@ -48,55 +48,18 @@ class FlashinferGemmaRmsnorm(OpWrapper): why="MTP pre_e_norm: the contiguous [T, 2560] embedding", spec=dict(shape=(4, 2560), dtype=torch.bfloat16, eps=1e-6), ), - Cell( - why="hidden 5120, 1 token, contiguous", - spec=dict(shape=(1, 5120), dtype=torch.bfloat16, eps=1e-6), - ), - Cell( - why="hidden 5120, 33 tokens, contiguous", - spec=dict(shape=(33, 5120), dtype=torch.bfloat16, eps=1e-6), - ), Cell( why="hidden 5120, 8192 tokens: the largest row count measured", spec=dict(shape=(8192, 5120), dtype=torch.bfloat16, eps=1e-6), ), - Cell( - why="trained-like weight ~0.02: the +1 dominates the scale", - spec=dict(shape=(4, 5120), dtype=torch.bfloat16, eps=1e-6, weight_scale=0.02), - ), - Cell(why="fp16", spec=dict(shape=(4, 5120), dtype=torch.float16, eps=1e-6)), - Cell(why="fp32", spec=dict(shape=(4, 5120), dtype=torch.float32, eps=1e-6)), - Cell( - why="3-D [16, 32, 128]: the per-head kernel", - spec=dict(shape=(16, 32, 128), dtype=torch.bfloat16, eps=1e-5), - ), - Cell( - why="column slice buf[:, :5120] of a [8, 10240] buffer", - spec=dict(shape=(8, 5120), col_offset=0, **_STREAM), - ), - Cell( - why="row slice rows[::2]: row stride 2 x width", - spec=dict(shape=(4, 5120), row_step=2, dtype=torch.bfloat16, eps=1e-6), - ), - Cell( - why="hidden 111: odd, not a multiple of the 8-element bf16 vector", - spec=dict(shape=(16, 111), dtype=torch.bfloat16, eps=1e-6), - ), - Cell( - why="hidden 1152: not a power of two", - spec=dict(shape=(16, 1152), dtype=torch.bfloat16, eps=1e-6), - ), - Cell( - why="hidden 6144", - spec=dict(shape=(16, 6144), dtype=torch.bfloat16, eps=1e-6), - ), ) note = """ Metadata consumed: none. Registered only when flashinfer is importable. - Certified path: flashinfer's CuTe DSL kernels (`rmsnorm_cute` for 2-D, - `qk_rmsnorm_cute` for 3-D, `weight_bias=1.0`); a strided row selects a - strided kernel variant. The CUDA-JIT fallback (`FLASHINFER_USE_CUDA_NORM=1`, + Certified path: flashinfer's CuTe DSL `rmsnorm_cute` kernel for 2-D bf16 + `x` (`weight_bias=1.0`), contiguous and at the residual's 10240 row stride; + a strided row selects a strided kernel variant. The 3-D per-head kernel + (`qk_rmsnorm_cute`) and fp16/fp32 operands are not certified. 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. @@ -111,7 +74,7 @@ class FlashinferGemmaRmsnorm(OpWrapper): width (8 elements for bf16/fp16, 4 for fp32) and assumes 16-byte aligned pointers, so odd row strides and misaligned slices were not measured. - Numerics: default tolerance holds for all dtypes, nothing loosened; the + Numerics: default tolerance holds for bf16, 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. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/layer_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/layer_norm.py index 84c4deb978a8..b8fe44d648e1 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/layer_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/layer_norm.py @@ -41,50 +41,19 @@ class LayerNorm(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec: shape, dtype; `param_dtype` (default dtype) for weight and bias; - # `bias` (default True); `layout` "contiguous" (default) or "column_slice" - # (the second half of a [rows, 2N] per-token buffer). Weight and bias are - # drawn over the checkpoint's measured ranges. - CELLS: tuple[Cell, ...] = ( - tuple( - Cell( - why=f"vision norm, {rows} patch rows x 1152 bf16", - spec=dict(shape=(rows, WIDTH), dtype=torch.bfloat16), - ) - for rows in (1, 7, 256, 1024, 4096) - ) - + ( - Cell( - why="vision norm at the tower call ceiling, 32768 patch rows x 1152 bf16", - spec=dict(shape=(32768, WIDTH), dtype=torch.bfloat16), - ), - ) - + tuple( - Cell( - why=f"{str(dtype)[6:]} activations, {rows} rows", - spec=dict(shape=(rows, WIDTH), dtype=dtype), - ) - for dtype in (torch.float16, torch.float32) - for rows in (7, 1024) - ) - + ( - Cell( - why="fp32 weight/bias against bf16 x: widened on load", - spec=dict(shape=(512, WIDTH), dtype=torch.bfloat16, param_dtype=torch.float32), - ), - Cell( - why="bias=None drops the additive term", - spec=dict(shape=(256, WIDTH), dtype=torch.bfloat16, bias=False), - ), - Cell( - why="3-D input keeps its shape", - spec=dict(shape=(4, 64, WIDTH), dtype=torch.bfloat16), - ), - Cell( - why="column slice of a [rows, 2N] buffer, read in place", - spec=dict(shape=(512, WIDTH), dtype=torch.bfloat16, layout="column_slice"), - ), + # spec: shape, dtype (weight and bias share it). Weight and bias are drawn + # over the checkpoint's measured ranges. + CELLS: tuple[Cell, ...] = tuple( + Cell( + why=f"vision norm, {rows} patch rows x 1152 bf16", + spec=dict(shape=(rows, WIDTH), dtype=torch.bfloat16), ) + for rows in (1, 1024) + ) + ( + Cell( + why="vision norm at the tower call ceiling, 32768 patch rows x 1152 bf16", + spec=dict(shape=(32768, WIDTH), dtype=torch.bfloat16), + ), ) note = """ @@ -93,7 +62,11 @@ class LayerNorm(OpWrapper): Layout: `x.stride(-1) == 1` is read in place through its row stride; any other `x`, and non-contiguous `weight`/`bias`, are silently `.contiguous()`d - first (one extra copy, correct result). `weight`/`bias` dtype is free. + first (one extra copy). `weight`/`bias` dtype is free. + + Certified: contiguous bf16 `[rows, 1152]` `x` with a bf16 weight and bias, + as the vision tower calls it. Not certified: strided or 3-D `x`, + `bias=None`, fp16/fp32 activations, weight/bias in another dtype than `x`. Loud failures (left to the callable): weight or bias length != N raises a bare AssertionError with no message; `next_pow2(N) * x.element_size() > diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/rms_norm_gated_token_major.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/rms_norm_gated_token_major.py index 0cce6669186d..a15d7599886e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/rms_norm_gated_token_major.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/rms_norm_gated_token_major.py @@ -80,21 +80,6 @@ class RmsNormGatedTokenMajor(OpWrapper): ARCHS = frozenset({Arch.SM_103}) CELLS: tuple[Cell, ...] = ( - Cell(why="decode: 1 token, checkpoint GDN epilogue, sigmoid", spec=dict(tokens=1, **_CKPT)), - Cell(why="17 tokens, checkpoint geometry, sigmoid", spec=dict(tokens=17, **_CKPT)), - Cell( - why="prefill 2048 tokens, checkpoint geometry, sigmoid", spec=dict(tokens=2048, **_CKPT) - ), - Cell(why="silu gate (the op's default), 2 tokens", spec=dict(tokens=2, **_SILU)), - Cell(why="silu gate, 512 tokens", spec=dict(tokens=512, **_SILU)), - Cell( - why="z a column slice of a 16384-wide projection, 1 token", - spec=dict(tokens=1, z="column_slice", **_SILU), - ), - Cell( - why="z a column slice: token stride 16384, 512 tokens", - spec=dict(tokens=512, z="column_slice", **_CKPT), - ), *( Cell( why=f"z the in_proj column view: token stride 16480, {t} tokens", @@ -106,39 +91,7 @@ class RmsNormGatedTokenMajor(OpWrapper): why="683 tokens: the first launch at 8 rows a program", spec=dict(tokens=683, z="in_proj", **_CKPT), ), - Cell( - why="1001 tokens x 48 heads at 8 rows a program, silu", - spec=dict(tokens=1001, z="in_proj", **_SILU), - ), - 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 = """ @@ -160,12 +113,16 @@ class RmsNormGatedTokenMajor(OpWrapper): extra elements (only `[0]` read) or not fp32 (widened silently). `weight` and `z` are widened to fp32 and need not match `x.dtype`. The - output is freshly allocated and contiguous even for a strided `x`; no input - is mutated. + output is freshly allocated and contiguous; no input is mutated. + + Certified: bf16 `x` with the sigmoid gate over the in_proj column view of + `z` (48 heads of N = 128) on the 4-row and 8-row kernels, plus one fp8-output + cell (silu gate, packed `z`). Not certified: the generic fallback kernel, + the silu gate without fp8 output, fp16/fp32 operands, a strided `x`. - Numerics: bf16 bitwise identical to the fp32 reference up to 816 rows; at - 24576 and 196608 rows one-ulp differences on ~1e-5 of elements, inside the - default tolerance. fp8 output byte-identical. Nothing loosened. + Numerics: bf16 inside the default tolerance against the fp32 reference, + with one-ulp differences on ~1e-5 of elements at prefill row counts. fp8 + output byte-identical. Nothing loosened. Hard limits N * itemsize fallback holds a row in 64 KB boundary 65536 B diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/gated_ngram_injection.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/gated_ngram_injection.py index 7cb589e7c5a1..cded9afd1790 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/gated_ngram_injection.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/gated_ngram_injection.py @@ -68,59 +68,28 @@ class GatedNgramInjection(OpWrapper): ARCHS = frozenset({Arch.SM_103}) # spec: kind "ckpt" (default) -- operands with this checkpoint's norm-gamma - # and value statistics: leading, hc_count, dtype, value_dtype (= dtype), - # clamp_min. kind "single_channel" -- one nonzero channel per stream so each - # dot is an exact product under the test's control: dots, value_dtype, - # clamp_min. + # and value statistics: leading, value_dtype (bf16), clamp_min. kind + # "single_channel" -- one nonzero channel per stream so each dot is an exact + # product under the test's control: dots, value_dtype, clamp_min. CELLS: tuple[Cell, ...] = ( - tuple( - Cell(why=f"checkpoint 4x2560 bf16, {why}", spec=dict(leading=(t,))) - for t, why in ( - (1, "decode, 1 token"), - (2, "2 tokens"), - (8, "decode batch 8"), - (33, "33 tokens / MTP verify 8 x 4 + 1"), - (512, "prefill 512 tokens"), - (4096, "prefill 4096 tokens"), - (8192, "prefill 8192 tokens, largest measured"), - ) - ) - + (Cell(why="[3, 17, width] leading dims", spec=dict(leading=(3, 17))),) - + tuple( - Cell(why=f"hc_count {c} at D 2560, 64 tokens", spec=dict(leading=(64,), hc_count=c)) - for c in (1, 2, 8) - ) - + tuple( - Cell( - why=f"{str(d).removeprefix('torch.')} operands, 512 tokens", - spec=dict(leading=(512,), dtype=d), - ) - for d in (torch.float16, torch.float32) - ) - + ( - Cell( - why="bf16 key/query with an fp32 value returns fp32", - spec=dict(leading=(128,), value_dtype=torch.float32), + Cell(why="checkpoint 4x2560 bf16, decode, 1 token", spec=dict(leading=(1,))), + Cell( + why="checkpoint 4x2560 bf16, 33 tokens / MTP verify 8 x 4 + 1", + spec=dict(leading=(33,)), + ), + Cell( + why="checkpoint 4x2560 bf16, prefill 8192 tokens, largest measured", + spec=dict(leading=(8192,)), + ), + Cell( + why="dots 0, +-5e-8 (inside the floor), 0.05, fp32 value", + spec=dict( + kind="single_channel", + dots=(0.0, 5e-8, -5e-8, 0.05), + value_dtype=torch.float32, + clamp_min=CKPT_CLAMP_MIN, ), - Cell( - why="dots 0, +-5e-8 (inside the floor), 0.05, fp32 value", - spec=dict( - kind="single_channel", - dots=(0.0, 5e-8, -5e-8, 0.05), - value_dtype=torch.float32, - clamp_min=CKPT_CLAMP_MIN, - ), - ), - Cell( - why="same dots with clamp_min 0.0, fp32 value", - spec=dict( - kind="single_channel", - dots=(0.0, 5e-8, -5e-8, 0.05), - value_dtype=torch.float32, - clamp_min=0.0, - ), - ), - ) + ), ) note = """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ngram_token_history.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ngram_token_history.py index 50c9d72d4924..704fc63f6b14 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ngram_token_history.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ngram_token_history.py @@ -88,173 +88,55 @@ class NgramTokenHistory(OpWrapper): ARCHS = frozenset({Arch.SM_103}) # spec: lengths, slots, carries, pool_slots (the cache manager's width, not - # the batch's), ngram_size (3), eos_density (0.0, of the packed ids), - # ids_dtype (int64), index_dtype (int32, for state_indices), qsl_dtype - # (int32), qsl_from "literal" | "seq_lens" (zeros + int64 cumsum of an - # int32 seq_lens prefix -- `cu_seqlens_from_seq_lens`' form), layout - # "contiguous" | "strided_ids" | "pool_view" (`pool[::2]`). + # the batch's), eos_density (0.0, of the packed ids), ids_dtype (int64), + # index_dtype (int32, for state_indices), qsl_dtype (int32), qsl_from + # "literal" | "seq_lens" (zeros + int64 cumsum of an int32 seq_lens prefix -- + # `cu_seqlens_from_seq_lens`' form). `ngram_size` is the checkpoint's 3. CELLS: tuple[Cell, ...] = ( - tuple( - _cell( - f"decode {n} x 1 token, pool sized to the batch", - (1,) * n, - range(n), - (True,) * n, - n + 2, - ) - for n in (1, 8, 32) - ) - + tuple( - _cell( - f"decode {n} x 1 token at the 461-slot engine pool", - (1,) * n, - _scattered(n), - (True,) * n, - ENGINE_SLOTS, - **_CALL_SITE, - ) - for n in (1, 8, 32) - ) - + tuple( - _cell( - f"prefill {'+'.join(map(str, ls))}, fresh, EOS in stream", - ls, - range(len(ls)), - (False,) * len(ls), - len(ls) + 2, - eos_density=0.05, - ) - for ls in ((37,), (1024,), (37, 1, 128), (7, 7, 7, 7)) - ) - + ( - _cell( - "mixed 64+13 contexts + 6 decodes, 10-slot pool", - (64, 13) + (1,) * 6, - (4, 0, 7, 2, 5, 1, 6, 3), - (False,) + (True,) * 7, - 10, - eos_density=0.03, - ), - _cell( - "mixed 64+13 contexts + 6 decodes at 461 slots", - (64, 13) + (1,) * 6, - (0, ENGINE_SLOTS - 1, 4, 17, 230, 231, 459, 1), - (False,) + (True,) * 7, - ENGINE_SLOTS, - eos_density=0.03, - **_CALL_SITE, - ), - _cell( - "MTP verify p=2: 8 decodes x 2 rows, qsl rebuilt", - (2,) * 8, - _scattered(8), - (True,) * 8, - ENGINE_SLOTS, - qsl_from="seq_lens", - **_CALL_SITE, - ), - _cell( - "MTP verify p=4: 32 decodes x 4 rows, qsl rebuilt", - (4,) * 32, - _scattered(32), - (True,) * 32, - ENGINE_SLOTS, - qsl_from="seq_lens", - **_CALL_SITE, - ), - *( - _cell( - f"MTP verify p={rows}: {decodes} decode{'s' if decodes > 1 else ''} x " - f"{rows} rows, qsl rebuilt", - (rows,) * decodes, - _scattered(decodes), - (True,) * decodes, - ENGINE_SLOTS, - qsl_from="seq_lens", - **_CALL_SITE, - ) - for rows, decodes in ((5, 8), (6, 32), (7, 8), (8, 1), (8, 32)) - ), - _cell( - "MTP verify mixed: 37-token context + 8 x 3 rows", - (37,) + (3,) * 8, - _scattered(9), - (False,) + (True,) * 8, - ENGINE_SLOTS, - qsl_from="seq_lens", - eos_density=0.05, - **_CALL_SITE, - ), - _cell( - "fresh sequences seed EOS, never the stale slot", (3, 5), (2, 0), (False, False), 4 - ), - _cell("short sequence publishes part carry, part new", (1,), (2,), (True,), 5), - _cell( - "zero-token rows: carry kept, fresh reseeded", - (2, 0, 0), - (0, 1, 2), - (True, True, False), - 3, - ), - _cell( - "unsorted slots, carry flags mixed, 8-slot pool", - (4, 6, 2), - (5, 0, 3), - (True, False, True), - 8, - ), - _cell("unnamed slots untouched, 8-slot pool", (5, 3), (1, 4), (True, True), 8), - _cell( - "in-stream EOS passed through raw", - (17, 1, 5, 1), - (3, 0, 2, 1), - (False, True, True, True), - 6, - eos_density=0.3, - ), - ) - + tuple( - _cell( - f"ids {str(i)[6:]}, slots/qsl {str(x)[6:]}", - (9, 1, 4), - (2, 0, 1), - (True, False, True), - 5, - ids_dtype=i, - index_dtype=x, - qsl_dtype=x, - ) - for i in (torch.int32, torch.int64) - for x in (torch.int32, torch.int64) - ) - + ( - _cell( - "ngram_size 4: a 3-id carry row", - (5, 1, 12), - (0, 2, 1), - (True, True, False), - 5, - ngram_size=4, - eos_density=0.05, - ), - _cell("ngram_size 2: a 1-id carry row", (6, 1), (1, 0), (False, True), 4, ngram_size=2), - _cell( - "input_ids a stride-2 view", - (9, 1, 4), - (2, 0, 1), - (True, False, True), - 5, - layout="strided_ids", - ), - _cell( - "context_pool a pool[::2] slot view, written through", - (9, 1, 4), - (2, 0, 1), - (True, False, True), - 5, - layout="pool_view", - ), - ) + _cell( + "decode 32 x 1 token at the 461-slot engine pool", + (1,) * 32, + _scattered(32), + (True,) * 32, + ENGINE_SLOTS, + **_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=8: 32 decodes x 8 rows, qsl rebuilt", + (8,) * 32, + _scattered(32), + (True,) * 32, + ENGINE_SLOTS, + qsl_from="seq_lens", + **_CALL_SITE, + ), + _cell( + "MTP verify mixed: 37-token context + 8 x 3 rows", + (37,) + (3,) * 8, + _scattered(9), + (False,) + (True,) * 8, + ENGINE_SLOTS, + qsl_from="seq_lens", + eos_density=0.05, + **_CALL_SITE, + ), + _cell("prefill 1024, fresh, EOS in stream", (1024,), (0,), (False,), 3, eos_density=0.05), + _cell( + "zero-token rows: carry kept, fresh reseeded", + (2, 0, 0), + (0, 1, 2), + (True, True, False), + 3, + ), ) note = """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_gate_conv_inject_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_gate_conv_inject_norm.py index 9a7ef343b849..70e3f55490db 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_gate_conv_inject_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_gate_conv_inject_norm.py @@ -344,43 +344,21 @@ class PleGateConvInjectNorm(OpWrapper): ARCHS = frozenset({Arch.SM_103}) # spec: lengths (rows per sequence), pool_slots (default num_seqs + 2), - # slots (None -> scattered distinct), max_rows (max(lengths)), pad (index - # of a sequence to give pad_slot_id), hc_count (4), hidden (2560), dtype - # (bf16), layout "contiguous" | "row_slice" (hc and key are buf[::2]). - CELLS: tuple[Cell, ...] = tuple( - _cell(f"decode {n} x 1 row at the 461-slot pool", (1,) * n, pool_slots=ENGINE_SLOTS) - for n in (1, 8, 32, 128) - ) + ( - _cell("decode 256 x 1 row, 512-slot pool", (1,) * 256, pool_slots=512), - _cell("MTP verify p=2: 8 x 2 rows", (2,) * 8, pool_slots=ENGINE_SLOTS), + # max_rows (max(lengths)), carries (None -> the argument is None). The + # checkpoint's 4 x 2560 bf16 geometry throughout. + CELLS: tuple[Cell, ...] = ( + _cell("decode 1 x 1 row at the 461-slot pool", (1,), pool_slots=ENGINE_SLOTS), + _cell("decode 128 x 1 row at the 461-slot pool", (1,) * 128, pool_slots=ENGINE_SLOTS), _cell("MTP verify p=4: 32 x 4 rows", (4,) * 32, pool_slots=ENGINE_SLOTS), - _cell("MTP verify p=3: 4 x 3 rows", (3,) * 4), - _cell("a pad_slot_id row beside live rows", (1,) * 6, pad=2), _cell("rows shorter than max_rows, a zero-row sequence", (2, 0, 1), max_rows=4), - _cell("hc and key row slices buf[::2]", (1,) * 16, layout="row_slice"), - _cell("hc_count 2 at hidden 2560", (1,) * 8, hc_count=2), - _cell("hidden 1024 (not the checkpoint)", (1,) * 8, hidden=1024), # Context-shaped (prefill / mixed) steps: the packed-row conv. _cell("prefill 8192 rows, fresh", (8192,), carries=(False,), pool_slots=ENGINE_SLOTS), - _cell( - "chunked prefill 16384 rows carrying", - (16384,), - carries=(True,), - pool_slots=ENGINE_SLOTS, - ), _cell( "mixed step: 2 contexts + 64 decode rows", (3001, 5000) + (1,) * 64, carries=(False, True) + (True,) * 64, pool_slots=ENGINE_SLOTS, ), - _cell( - "short contexts, tile-straddling, zero-row, pad", - (40, 1, 0, 7, 33, 2), - carries=(True, False, True, False, True, True), - pad=4, - ), - _cell("mixed step, carry None", (9, 1, 1)), ) note = """ @@ -394,7 +372,8 @@ class PleGateConvInjectNorm(OpWrapper): Preconditions (data; not guarded): at most `max_rows` rows per sequence; distinct live slots; `query_start_loc[-1] == T`. Operands 2-D with unit - column stride (row strides free); weights contiguous. + column stride; weights contiguous. Certified at the checkpoint's 4 x 2560 + bf16 geometry with every slot live (the target never passes `pad_slot_id`). Numerics: every rounding point of the composition is kept, so the only differences are fp32 reduction order (the three norms' sums of squares diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_ngram_embed.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_ngram_embed.py index 27d3a5e80837..22f3182ff674 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_ngram_embed.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_ngram_embed.py @@ -309,77 +309,48 @@ class PleNgramEmbed(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec: lengths (rows per sequence), slots (None -> arange), pool_slots - # (default num_seqs + 2), carries (None -> the argument is None), max_rows - # (default max(lengths)), eos_density (0.0), table "small" | "host" | - # "ckpt" (the full 320001536-row device table, rows filled where touched), - # ngram_size (3), heads_per_ngram (8), ids_dtype (int32). + # spec: lengths (rows per sequence), pool_slots (default num_seqs + 2), + # carries (None -> the argument is None), max_rows (default max(lengths)), + # eos_density (0.0), table "small" | "host" | "ckpt" (the full + # 320001536-row device table, rows filled where touched). The checkpoint's + # 3-gram of 8 heads per order, int32 ids. CELLS: tuple[Cell, ...] = ( - tuple( - _cell(f"decode {n} x 1 row at the 461-slot pool", (1,) * n, pool_slots=ENGINE_SLOTS) - for n in (1, 8, 32, 128) - ) - + tuple( - _cell( - f"decode {n} x 1 row, host-resident table, 461 slots", - (1,) * n, - pool_slots=ENGINE_SLOTS, - table="host", - ) - for n in (1, 128) - ) - + ( - _cell( - "decode 128 x 1 row, the full checkpoint-height table", - (1,) * 128, - pool_slots=ENGINE_SLOTS, - table="ckpt", - eos_density=0.05, - ), - _cell("decode 64 x 1 row, EOS in stream and carry", (1,) * 64, eos_density=0.3), - _cell("MTP verify p=2: 8 x 2 rows", (2,) * 8, pool_slots=ENGINE_SLOTS), - _cell("MTP verify p=3: 8 x 3 rows", (3,) * 8, eos_density=0.2), - _cell("MTP verify p=4: 32 x 4 rows", (4,) * 32, pool_slots=ENGINE_SLOTS), - _cell( - "carry flags mixed, fresh rows seed EOS", - (1, 3, 1, 2), - carries=(True, False, True, False), - ), - _cell( - "zero-row sequences republish carry", - (2, 0, 1, 0), - carries=(True, True, True, False), - max_rows=2, - ), - _cell("rows shorter than max_rows", (1, 3, 2), max_rows=4), - _cell("int64 ids", (1,) * 16, ids_dtype=torch.int64), - _cell("ngram_size 4, 4 heads per order", (1, 2, 4), ngram_size=4, heads_per_ngram=4), - _cell("ngram_size 2, 8 heads", (1,) * 8, ngram_size=2, eos_density=0.2), - # Context-shaped (prefill / mixed) steps: the packed-row form. - _cell("prefill 8192 rows, fresh", (8192,), carries=(False,), pool_slots=ENGINE_SLOTS), - _cell( - "chunked prefill 16384 rows carrying, host table", - (16384,), - carries=(True,), - pool_slots=ENGINE_SLOTS, - table="host", - eos_density=0.01, - ), - _cell( - "mixed step: 2 contexts + 64 decode rows", - (3001, 5000) + (1,) * 64, - carries=(False, True) + (True,) * 64, - pool_slots=ENGINE_SLOTS, - eos_density=0.05, - ), - _cell( - "short contexts, tile-straddling, zero-row", - (40, 1, 0, 7, 33, 2), - carries=(True, False, True, False, True, True), - eos_density=0.2, - ), - _cell("mixed step, carry None", (9, 1, 1), max_rows=9), - ) + _cell( + "decode 128 x 1 row, host-resident table, 461 slots", + (1,) * 128, + pool_slots=ENGINE_SLOTS, + table="host", + ), + _cell( + "decode 128 x 1 row, the full checkpoint-height table", + (1,) * 128, + pool_slots=ENGINE_SLOTS, + table="ckpt", + eos_density=0.05, + ), + _cell("MTP verify p=4: 32 x 4 rows", (4,) * 32, pool_slots=ENGINE_SLOTS), + _cell( + "zero-row sequences republish carry", + (2, 0, 1, 0), + carries=(True, True, True, False), + max_rows=2, + ), + # Context-shaped (prefill / mixed) steps: the packed-row form. + _cell( + "chunked prefill 16384 rows carrying, host table", + (16384,), + carries=(True,), + pool_slots=ENGINE_SLOTS, + table="host", + eos_density=0.01, + ), + _cell( + "mixed step: 2 contexts + 64 decode rows", + (3001, 5000) + (1,) * 64, + carries=(False, True) + (True,) * 64, + pool_slots=ENGINE_SLOTS, + eos_density=0.05, + ), ) note = """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/bf16_to_fp32.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/bf16_to_fp32.py index d34139053140..fae7a1dbaaea 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/bf16_to_fp32.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/bf16_to_fp32.py @@ -43,14 +43,11 @@ class Bf16ToFp32(OpWrapper): CELLS: tuple[Cell, ...] = ( Cell(why="decode logits: 128 rows of the 248320 vocabulary", spec=dict(shape=(128, VOCAB))), - Cell(why="mixed step logits: 124 rows", spec=dict(shape=(124, VOCAB))), - Cell(why="one row", spec=dict(shape=(1, VOCAB))), + Cell(why="one row; a ragged last program", spec=dict(shape=(1, VOCAB))), Cell( why="every bf16 bit pattern: NaNs, infinities, subnormals, signed zeros", spec=dict(shape=(1, 65536), source="all_codes"), ), - Cell(why="3 x 7: below one program", spec=dict(shape=(3, 7))), - Cell(why="rank 3, a ragged last program", spec=dict(shape=(2, 5, 4099))), Cell(why="zero rows", spec=dict(shape=(0, VOCAB))), ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/dequantize_e4m3_per_tensor.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/dequantize_e4m3_per_tensor.py index 68888e530047..dd2eb76c77a4 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/dequantize_e4m3_per_tensor.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/dequantize_e4m3_per_tensor.py @@ -43,66 +43,15 @@ class DequantizeE4m3PerTensor(OpWrapper): CELLS: tuple[Cell, ...] = ( Cell(why="_ple_block decode: [1, 16, 160] gathered rows, bf16 scale", spec=_ple(1)), Cell(why="_ple_block at 8 tokens", spec=_ple(8)), - Cell(why="_ple_block prefill: [4096, 16, 160]", spec=_ple(4096)), Cell( - why="every e4m3 byte into bf16, five scale values, NaN codes included", - spec=dict( - shape=(1, 256), dtype=torch.bfloat16, scales=_DOMAIN_SCALES, source="all_codes" - ), - ), - Cell( - why="every e4m3 byte into fp16", - spec=dict( - shape=(1, 256), dtype=torch.float16, scales=_DOMAIN_SCALES, source="all_codes" - ), - ), - Cell( - why="every e4m3 byte into fp32", - spec=dict( - shape=(1, 256), dtype=torch.float32, scales=_DOMAIN_SCALES, source="all_codes" - ), - ), - Cell( - why="rank 2: the scale's rank tracks the input's", - spec=dict(shape=(4, 8), dtype=torch.bfloat16, scales=(PLE_SCALE,), source="randn"), - ), - Cell( - why="rank 5: the deepest rank run", - spec=dict( - shape=(2, 2, 2, 4, 8), dtype=torch.bfloat16, scales=(PLE_SCALE,), source="randn" - ), - ), - Cell( - why="1 x 1: far below the fixed 1024-block grid", - spec=dict(shape=(1, 1), dtype=torch.bfloat16, scales=(PLE_SCALE,), source="randn"), - ), - Cell( - why="3 x 7: odd shape, no width or alignment rule", - spec=dict(shape=(3, 7), dtype=torch.bfloat16, scales=(PLE_SCALE,), source="randn"), - ), - Cell( - why="[2^20, 160]: 167,772,160 elements through the grid-stride loop", - spec=dict( - shape=(1 << 20, 160), dtype=torch.bfloat16, scales=(PLE_SCALE,), source="randn" - ), - ), - Cell( - why="scale as a one-element slice at a nonzero storage offset", - spec=dict( - shape=(8, 160), - dtype=torch.bfloat16, - scales=(PLE_SCALE,), - source="randn", - scale_slice=True, - ), + why="_ple_block prefill at max_num_tokens: [16384, 16, 160] through the grid-stride " + "loop", + spec=_ple(16384), ), Cell( - why="input at a non-16-byte-aligned address: same kernel", + why="every e4m3 byte into bf16, five scale values, NaN codes included", spec=dict( - shape=(64, 160), - dtype=torch.bfloat16, - scales=(PLE_SCALE, 1 / 3, 7.3e-3), - source="unaligned", + shape=(1, 256), dtype=torch.bfloat16, scales=_DOMAIN_SCALES, source="all_codes" ), ), ) @@ -129,9 +78,11 @@ class DequantizeE4m3PerTensor(OpWrapper): `scales` value written in place (the scale is read from memory at replay). Hard limits - numel probed to 167,772,160; the int32 boundary 2^31 not probed + numel certified to 41,943,040 (16384 tokens); probed to 167,772,160; the + int32 boundary 2^31 not probed - Not certified: `TRTLLM_ENABLE_PDL=1`, `scales` on a different device than + Not certified: fp16 and fp32 outputs (the target's scale is bf16), ranks + other than 3, `TRTLLM_ENABLE_PDL=1`, `scales` on a different device than `input`. """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/fp8_quantize_1x128.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/fp8_quantize_1x128.py index 8129463ba804..1472661e7d05 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/fp8_quantize_1x128.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/fp8_quantize_1x128.py @@ -72,20 +72,13 @@ class Fp8Quantize1x128(OpWrapper): CELLS: tuple[Cell, ...] = ( Cell(why="MTPLayer._moe decode: one token at hidden 2560", spec=dict(tokens=1, **_CKPT)), Cell( - why="3 tokens: inside the data buffer's 4-row allocation pad", - spec=dict(tokens=3, **_CKPT), - ), - Cell(why="4 tokens: exactly one 4-row pad unit", spec=dict(tokens=4, **_CKPT)), - Cell( - why="5 tokens: the first count spilling into a second pad unit", + why="5 tokens: the first count spilling into a second 4-row 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 of the identity config", spec=dict(tokens=8192, **_CKPT), @@ -94,29 +87,6 @@ class Fp8Quantize1x128(OpWrapper): why="16384 rows: max_num_tokens of the tuned config", spec=dict(tokens=16384, **_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 = """ @@ -144,7 +114,8 @@ class Fp8Quantize1x128(OpWrapper): Hard limits T certified 0..16384 at H 2560 (max_num_tokens of the tuned config); above that not probed - H multiple of 16 (raises otherwise); certified 16..7168 + H multiple of 16 (raises otherwise); certified at 2560, the draft layer's + width (the test also drives a 16-wide tail block) """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_fwd.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_fwd.py index d1932d41571f..fb73e8eda6bc 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_fwd.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_fwd.py @@ -63,172 +63,54 @@ class CausalConv1dFwd(OpWrapper): # 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="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="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="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="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="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="channel-last x (a bare .t() of token-major), into its own out", - spec=dict( - dim=CKPT_DIM, - lengths=(37, 1, 128), - pool_slots=8, - slots=(3, 0, 5), - initial=(True, False, True), - layout="channel_last", - ), - ), - Cell( - why="pad_slot_id -1 (the engine's sentinel): sequence and slot skipped", - spec=dict( - dim=1024, - lengths=(16, 8, 16), - pool_slots=4, - slots=(0, -1, 2), - initial=(True, True, True), - ), - ), - Cell( - why="pad_slot_id 99: the sentinel is compared by value", - spec=dict( - dim=1024, - lengths=(16, 8, 16), - pool_slots=4, - slots=(0, 99, 2), - initial=(True, True, True), - pad_slot_id=99, - ), - ), - Cell( - why="sequences shorter than the window: state part carry, part tokens", - spec=dict( - dim=512, - lengths=(2, 1, 2), - pool_slots=4, - slots=(0, 1, 2), - initial=(True, True, False), - ), - ), - Cell( - why="a zero-token sequence republishes its carry unchanged", - spec=dict( - dim=512, lengths=(0, 40), pool_slots=4, slots=(2, 3), initial=(True, False) - ), - ), - ) - + tuple( - Cell( - why=f"{str(dtype).removeprefix('torch.')} with bias, silu={silu}", - spec=dict( - dim=256, - lengths=(30, 34), - pool_slots=2, - slots=(0, 1), - initial=(True, True), - dtype=dtype, - bias=True, - silu=silu, - ), - ) - for dtype in (torch.bfloat16, torch.float16, torch.float32) - for silu in (True, False) - ) - + ( - Cell( - why="conv_states None: every sequence starts from zeros", - spec=dict(dim=512, lengths=(12, 20), pool_slots=None, slots=None, initial=None), - ), - Cell( - why="cache_indices/has_initial_state None: slot s, no carry", - spec=dict(dim=512, lengths=(24, 40), pool_slots=2, slots=None, initial=None), - ), - Cell( - why="dense [batch, dim, seqlen] layout, query_start_loc None", - spec=dict( - dim=256, batch3d=(3, 40), pool_slots=3, slots=None, initial=(True, False, True) - ), - ), - Cell( - why="one 2048-token sequence: the vectorised channel-major load", - spec=dict(dim=CKPT_DIM, lengths=(2048,), pool_slots=1, slots=(0,), initial=(True,)), - ), - Cell( - why="one 2047-token sequence: the scalar channel-major load", - spec=dict(dim=CKPT_DIM, lengths=(2047,), pool_slots=1, slots=(0,), initial=(True,)), - ), - Cell( - why="several long sequences: the 128-thread block", - spec=dict( - dim=1024, - lengths=(700, 701, 648), - pool_slots=3, - slots=(0, 1, 2), - initial=(True, True, True), - ), - ), - Cell( - why="every sequence a single token", - spec=dict( - dim=1024, - lengths=(1, 1, 1, 1), - pool_slots=4, - slots=(0, 1, 2, 3), - initial=(True, True, True, True), - ), - ), - ) - + tuple( - Cell( - why=f"conv width {width}", - spec=dict( - dim=512, - width=width, - lengths=(33, 7), - pool_slots=2, - slots=(0, 1), - initial=(True, True), - ), - ) - for width in (2, 3) - ) + ), + Cell( + why="one 2048-token sequence: the vectorised load, 128-thread block", + spec=dict(dim=CKPT_DIM, lengths=(2048,), pool_slots=1, slots=(0,), initial=(True,)), + ), + Cell( + why="one 2047-token sequence: the scalar load, 128-thread block", + spec=dict(dim=CKPT_DIM, lengths=(2047,), pool_slots=1, slots=(0,), initial=(True,)), + ), ) note = """ @@ -239,10 +121,12 @@ class CausalConv1dFwd(OpWrapper): Call site: the GDN prefill half -- `x` the `[10240, num_prefill_tokens]` channel-major output of `extract_transpose_prefill_slice`, W 4, no bias, SiLU on, the layer's whole conv pool unsliced, `pad_slot_id` -1, `out` None. - The first three cells are that call; 461 slots tops the certified pool - interval (8 and 461 bracket every width the cache manager was seen to - allocate). At 461 the activation and state are bitwise identical to the - 8-slot result for the same slot content. + The cells are that call (bf16, W 4, no bias, SiLU, channel-major in + place) at the engine's 461-slot pool, which tops every width the cache + manager was seen to allocate, and at small pools for the pad sentinel, + the sub-window sequences and both load forms of the 128-thread block + (single sequences of 2048 and 2047 tokens); the multi-sequence cells + run the 64-thread block with scalar loads. Preconditions the op checks itself (all raise): dtypes in {bf16, fp16, fp32}, `conv_states`/`bias` dtype, int32 `query_start_loc`/`cache_indices`, @@ -252,11 +136,14 @@ class CausalConv1dFwd(OpWrapper): Not checked, not guarded (content, not metadata, or an illegal access rather than a wrong number): a `cache_indices` value outside `[0, slots)`; `has_initial_state[s]` True with `conv_states=None` (reads through null); - `conv_states`' shape (only its dtype is validated). Sequences of length 0 - and shorter than `W - 1` are legal. Unnamed slots are bitwise untouched. + `conv_states`' shape (only its dtype is validated). Sequences shorter + than `W - 1` are legal and certified; a zero-length sequence, which the + target never schedules, is not certified. 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. + Only the channel-major kernel (the target's layout) is certified; the + channel-last kernel and `TRTLLM_CAUSAL_CONV1D_SILU_MODE`, which selects + faster SiLU forms on it, are not. Hard limits x token/channel strides uint32 stride fields boundary 2^32 elements diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_triton_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_triton_update.py index 73c73cf96afc..56437ad3cd92 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_triton_update.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_triton_update.py @@ -79,7 +79,7 @@ class CausalConv1dTritonUpdate(OpWrapper): CELLS: tuple[Cell, ...] = ( tuple( Cell( - why=f"verify at the 461-slot pool, batch {batch}, {seqlen} tokens/row", + why=f"verify, x over the [n*p, 16480] in_proj, batch {batch}, {seqlen} tok/row", spec=dict( dim=CKPT_DIM, batch=batch, @@ -87,39 +87,14 @@ class CausalConv1dTritonUpdate(OpWrapper): pool_slots=ENGINE_SLOTS, slots=_pinned(batch), snapshots=ENGINE_SNAPSHOT_ROWS, + x_form="in_proj", ), ) - for batch, seqlen in ( - (1, 2), - (8, 3), - (32, 2), - (17, 4), - # max_draft_len 4..7: seqlen 5..8, up to 256 rows. - (32, 5), - (8, 6), - (17, 7), - (1, 8), - (17, 8), - (32, 8), - ) - ) - + ( - Cell( - why="seqlen 5 into an 8-step snapshot buffer: steps 5..7 untouched", - spec=dict( - dim=CKPT_DIM, - batch=8, - seqlen=5, - pool_slots=ENGINE_SLOTS, - slots=_pinned(8), - snapshots=ENGINE_SNAPSHOT_ROWS, - snapshot_steps=8, - ), - ), + for batch, seqlen in ((1, 2), (32, 3)) ) + tuple( Cell( - why=f"verify, x over the [n*p, 16480] in_proj, batch {batch}, {seqlen} tok/row", + why=f"verify at the 461-slot pool, batch {batch}, {seqlen} tokens/row", spec=dict( dim=CKPT_DIM, batch=batch, @@ -127,130 +102,30 @@ class CausalConv1dTritonUpdate(OpWrapper): pool_slots=ENGINE_SLOTS, slots=_pinned(batch), snapshots=ENGINE_SNAPSHOT_ROWS, - x_form="in_proj", ), ) - for batch, seqlen in ((1, 2), (32, 3)) + # seqlen 8: max_draft_len 7, the widest verify step. + for batch, seqlen in ((17, 4), (32, 8)) ) - + tuple( + + ( Cell( - why=f"verify, 128-slot pool, batch {batch}, {seqlen} tokens/row", + why="seqlen 5 into an 8-step snapshot buffer: steps 5..7 untouched", spec=dict( dim=CKPT_DIM, - batch=batch, - seqlen=seqlen, - pool_slots=128, - slots="random", - snapshots=batch, + batch=8, + seqlen=5, + pool_slots=ENGINE_SLOTS, + slots=_pinned(8), + snapshots=ENGINE_SNAPSHOT_ROWS, + snapshot_steps=8, ), - ) - 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 - ), - ), ) ) @@ -263,17 +138,19 @@ class CausalConv1dTritonUpdate(OpWrapper): 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, seqlen 2..8 (a buffer deeper than seqlen keeps its - later steps untouched). + snapshots are what a rollback has to read. Certified at the 461-slot pool + with a 32-row snapshot buffer, `x` as the target's q|k|v column view of + the projection and as a contiguous `[n, p, 10240]` block, seqlen 2, 3, 4, + 5 and 8 (a buffer deeper than seqlen keeps its later steps untouched), and + for pad-sentinel rows; bf16, W 4, no bias, SiLU only. Raise loudly, so not guarded: width 1 (triton CompilationError); a missing `intermediate_state_indices` with snapshots and slot indices given. Not guarded (content): duplicate slots or duplicate snapshot rows race; - `num_accepted_tokens[b]` must lie in `[1, seqlen]` (certified 1..3 at - seqlen 3). Out-of-range slot values: illegal access, not probed. + `num_accepted_tokens[b]` must lie in `[1, seqlen]` (the accepted-token + roll is uncertified; the target passes None). Out-of-range slot values: + illegal access, not probed. Accepted and inert: `launch_dependent_kernels` (bitwise identical either way). Accepted but not exposed: `cache_seqlens` (ignored by the kernel -- diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_update.py index 2e1846e5d478..a56a5017fc83 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_update.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_update.py @@ -343,144 +343,27 @@ class CausalConv1dUpdate(OpWrapper): # pad_slot_id, x_form ("3d" | "2d" | "2d_row_slice" | "2d_in_proj" | # "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"x the q|k|v column view of the [rows, 16480] in_proj, batch {batch}", - spec=dict( - dim=CKPT_DIM, - batch=batch, - pool_slots=pool, - slots=slots, - x_form="2d_in_proj", - ), - ) - for batch, pool, slots in ( - (1, ENGINE_SLOTS, _pinned(1)), - (32, ENGINE_SLOTS, _pinned(32)), - (128, 256, "random"), - ) - ) - + 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), + CELLS: tuple[Cell, ...] = tuple( + Cell( + why=f"x the q|k|v column view of the [rows, 16480] in_proj, batch {batch}", + spec=dict( + dim=CKPT_DIM, + batch=batch, + pool_slots=pool, + slots=slots, + x_form="2d_in_proj", ), ) - + 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) + for batch, pool, slots in ( + (1, ENGINE_SLOTS, _pinned(1)), + (32, ENGINE_SLOTS, _pinned(32)), + (128, 256, "random"), ) + ) + ( + 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)), + ), ) note = """ @@ -492,19 +375,20 @@ class CausalConv1dUpdate(OpWrapper): the layer's whole conv pool `[slots, 10240, 3]`, W 4, no bias, SiLU on, `cache_seqlens` None, int32 slots, `pad_slot_id` -1. That exact shape (silu, no ring, `state_len == W - 1`) runs TRT-LLM's - dedicated seqlen-1 kernel; every other cell runs its general one. From - `_CUTE_MIN_ROWS` (64) rows, with bf16 operands, a contiguous weight, a - pool whose slot is `[dim, 3]` contiguous and 16-byte aligned rows, it runs + dedicated seqlen-1 kernel; any other shape runs its general one, which the + target never reaches and which is not certified. From `_CUTE_MIN_ROWS` + (64) rows, with bf16 operands, a contiguous weight, a pool whose slot is + `[dim, 3]` contiguous and 16-byte aligned rows, it runs `_DecodeConvKernel` (CuTe DSL) instead: 8 channels a thread by 16-byte loads, the TRT-LLM kernel's instructions, bitwise its result (x and pool, pad rows untouched); 128 rows 7.46 -> 5.41 us a layer (GB300, graph - replay over 8 rotating layers). Certified at 128 - and 461 pool slots; at 461 the result is bitwise identical to 128 for the - same slot content. + replay over 8 rotating layers). Certified at that call: 1 and 32 rows + over the 461-slot pool, 128 rows (the CuTe kernel) over a 256-slot pool, + and pad rows over a 128-slot pool. - 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 + The op checks metadata itself, loudly: width 2..4, `state_len >= W - 1`, + one dtype across x/conv_state/weight (bias matches weight), weight shape, + int32 contiguous length-batch `conv_state_indices`, int32 `cache_seqlens`, a batch-sized pool when indices are absent. So no guard. Not guarded because it is content, not metadata: two rows naming one slot @@ -512,9 +396,8 @@ class CausalConv1dUpdate(OpWrapper): 10 runs while both outputs matched. Slot uniqueness is the caller's. Out-of-range slot values are an illegal access, not probed. - A row whose index equals `pad_slot_id` is skipped: its x rows and the pool - stay bitwise unchanged (linear and ring mode). Any value works as the - sentinel, including one that collides with a live slot. + A row whose index equals `pad_slot_id` (-1 at the call site) is skipped: + its x rows and the pool stay bitwise unchanged. Hard limits none probed. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/chunk_gated_delta_rule.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/chunk_gated_delta_rule.py index 5f87469673ea..4a5d3c03d82c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/chunk_gated_delta_rule.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/chunk_gated_delta_rule.py @@ -73,102 +73,62 @@ class ChunkGatedDeltaRule(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec: seqlens (packed with cu_seqlens) or batch/num_tokens (no - # cu_seqlens); initial: "indexed" (a `slots`-row pool, `indices` or arange) - # | "positional" | None; pool_dtype (fp32 | bf16, the dtype of the indexed or - # positional pool); output: None | "buffer" | "window" (wide[:, 50:146] - # of a 200-token buffer); output_final_state; scale; normalize_qk (False - # with l2norm feeds raw q/k). + # spec: seqlens (packed with cu_seqlens); initial: "indexed" (a `slots`-row + # pool at `indices`) | "positional"; pool_dtype (fp32 | bf16, the dtype of + # the indexed pool); output: None | "buffer" | "window" (wide[:, 50:146] of + # a 200-token buffer). 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", - ), + Cell( + why="fp32 pool (a verify step's context rows), 461 slots, output= buffer", + spec=dict( + seqlens=(200, 37, 91), + initial="indexed", + slots=ENGINE_SLOTS, + indices=_engine_slots(3), + output="buffer", ), - ) - + tuple( - Cell( - why=f"bf16 pool, 461-slot pool, seqlens {list(s)}", - spec=dict( - seqlens=s, - initial="indexed", - slots=ENGINE_SLOTS, - indices=_engine_slots(len(s)), - pool_dtype=torch.bfloat16, - ), - ) - for s in ((1,), (64,), (200, 37, 91), (1024,), (200, 1, 1, 1, 1)) - ) - + ( - Cell( - why="bf16 pool, 461-slot pool, output= buffer", - spec=dict( - seqlens=(200, 37, 91), - initial="indexed", - slots=ENGINE_SLOTS, - indices=_engine_slots(3), - pool_dtype=torch.bfloat16, - output="buffer", - ), + ), + Cell( + why="fp32 pool, 461-slot pool, seqlens [1024]", + spec=dict( + seqlens=(1024,), initial="indexed", slots=ENGINE_SLOTS, indices=_engine_slots(1) ), - ) - + 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="bf16 pool, 461-slot pool, output= buffer", + spec=dict( + seqlens=(200, 37, 91), + initial="indexed", + slots=ENGINE_SLOTS, + indices=_engine_slots(3), + pool_dtype=torch.bfloat16, + output="buffer", ), - Cell( - why="no cu_seqlens: two equal-length rows", - spec=dict(batch=2, num_tokens=96, initial="positional", output_final_state=True), + ), + Cell( + why="bf16 pool, 461-slot pool, seqlens [1024]", + spec=dict( + seqlens=(1024,), + initial="indexed", + slots=ENGINE_SLOTS, + indices=_engine_slots(1), + pool_dtype=torch.bfloat16, ), - Cell( - why="explicit scale 0.05 replaces the default", - spec=dict(seqlens=(80,), initial=None, scale=0.05), + ), + Cell( + why="bf16 pool, mixed step: a context and four 1-token rows [200, 1, 1, 1, 1]", + spec=dict( + seqlens=(200, 1, 1, 1, 1), + initial="indexed", + slots=ENGINE_SLOTS, + indices=_engine_slots(5), + pool_dtype=torch.bfloat16, ), - Cell( - why="output= a token-window slice of a wider buffer", - spec=dict(seqlens=(96,), initial="positional", output="window"), - ), - ) + ), + Cell( + why="output= a token-window slice of a wider buffer", + spec=dict(seqlens=(96,), initial="positional", output="window"), + ), ) note = """ @@ -180,8 +140,7 @@ class ChunkGatedDeltaRule(OpWrapper): or bf16 as `mamba_ssm_cache_dtype` allocated it, int32 slots, int64 `cu_seqlens`; the MTP-verify context half passes the `si_p` / `qsl_long_p` prefixes, and a mixed batch's one-token decode rows ride in the same call. - Certified at small pools and at 461 slots; at 461 output and state are - bitwise identical to a 6-slot pool with the same slot content. + Certified at 461 slots. Over a bf16 pool the initial state is read as fp32 and the final state is stored rounded once to bf16, once per call and sequence; the accumulation is unchanged. @@ -190,8 +149,8 @@ class ChunkGatedDeltaRule(OpWrapper): warm one: the state kernel's autotuner snapshots and restores the indexed state slots its benchmark launches write (`chunk_delta_h.py`, `_save_autotune_state`). The key includes the pool dtype, so a bf16 pool has its own. Certified in a fresh interpreter - at this geometry over an fp32 and over a bf16 pool, and at 5 / 20 heads. The positional - path never aliased its state. + at this geometry over a bf16 pool (the fp32 pool's key takes the same snapshot and + restore). The positional path never aliased its state. Loud, from the callable: fp32 or mixed-dtype q/k/v, a 2-D `beta`, `q.shape[0] != 1` with `cu_seqlens`, a sequence-count mismatch of indices @@ -201,12 +160,14 @@ class ChunkGatedDeltaRule(OpWrapper): by the caller or the in-kernel flag (raw q/k with the flag off return finite values ~1e25). `head_first` is pinned False. - Sequence lengths are unconstrained: 1..1024, off and on the 64 boundary, - mixed ragged, all correct. Memory: a `[B, NT, H, K, V]` per-chunk state + Sequence lengths are unconstrained: certified at 1..1024, off and on the + 64 boundary, mixed ragged. Memory: a `[B, NT, H, K, V]` per-chunk state scratch in `q.dtype`, 24 MiB per 1024 tokens here. Deterministic. Not certified: the FlashInfer sibling, fp16 (accepted), bf16 `g`/`beta`, an - fp16 pool, `K = 256` (accepted, unmeasured), empty `T` or zero-length spans. + fp16 pool, `K = 256` (accepted, unmeasured), empty `T` or zero-length spans, + and what the target never passes: `initial_state=None`, `output_final_state`, + `use_qk_l2norm_in_kernel`, an explicit `scale`, no `cu_seqlens`. Hard limits K state kernel asserts K <= 256 (read, not probed) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/cu_seqlens_from_seq_lens.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/cu_seqlens_from_seq_lens.py index 843e1e9c78a2..9fb98404938a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/cu_seqlens_from_seq_lens.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/cu_seqlens_from_seq_lens.py @@ -20,18 +20,15 @@ #: 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..8: `max_draft_len` 1..7). +#: Batch shapes this checkpoint's serving path produces where it calls this entry. +#: ``p`` is the trunk's query rows per generation request, ``1 + runtime_draft_len`` +#: under MTP (2..8: `max_draft_len` 1..7); at ``p = 1`` the target reuses the +#: recurrent metadata's vector instead. BATCHES = { - "single_prefill": (4096,), - "decode_p1": (1,) * 32, "decode_p2": (2,) * 32, - "mixed_p1": (1024, 37, 1, 1, 1, 1), "mixed_p2": (1024, 37, 2, 2, 2, 2), "chunked_prefill_p2": (512, 512, 96, 2, 2, 2, 2, 2, 2, 2), "padded_decode_p2": (2, 2, 2, 0, 0, 0), - "one_sequence": (7,), "max_batch": (2,) * 256, } @@ -52,66 +49,32 @@ class CuSeqlensFromSeqLens(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec: lengths, in_dtype, out_dtype; `layout` "contiguous" (default), + # spec: lengths, in_dtype, out_dtype; `layout` "contiguous" (default) or # "prefix" (the first len(lengths) of a longer buffer -- the call site's - # `split(seq_lens_cuda, [batch, ...])[0]`), or "strided" (a [::2] view). + # `seq_lens_cuda[:batch]`). CELLS: tuple[Cell, ...] = ( - tuple( - Cell( - why=f"{name} {_name(i)}->{_name(o)}", - spec=dict(lengths=lengths, in_dtype=i, out_dtype=o), - ) - for name, lengths in BATCHES.items() - for i in (torch.int32, torch.int64) - for o in (torch.int32, torch.int64) - ) - + ( - Cell( - why="call site: prefix of seq_lens_cuda, int32->int64", - spec=dict( - lengths=BATCHES["mixed_p2"], - in_dtype=torch.int32, - out_dtype=torch.int64, - layout="prefix", - ), + Cell( + why="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=f"MTP verify p={p}, {requests} requests: prefix int32->int64", - spec=dict( - lengths=(p,) * requests, - in_dtype=torch.int32, - out_dtype=torch.int64, - layout="prefix", - ), - ) - for p, requests in ((5, 8), (6, 32), (7, 8), (8, 32)) + ), + Cell( + why="MTP verify p=8, 32 requests: prefix int32->int64", + spec=dict( + lengths=(8,) * 32, in_dtype=torch.int32, out_dtype=torch.int64, layout="prefix" ), + ), + *( Cell( - why="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), + why=f"{name} {_name(torch.int32)}->{_name(torch.int64)}", + spec=dict(lengths=BATCHES[name], in_dtype=torch.int32, out_dtype=torch.int64), ) - for d in (torch.uint8, torch.int16) - ) + for name in ("chunked_prefill_p2", "padded_decode_p2", "max_batch") + ), ) note = """ @@ -134,6 +97,9 @@ class CuSeqlensFromSeqLens(OpWrapper): a captured graph follows lengths rewritten in place). Three launches, about 0.03 ms at batch 8..256 -- launch bound. + Not certified: a `seq_lens` dtype other than int32 (the guard admits any + integer dtype), a non-contiguous `seq_lens` view. + Hard limits out[-1] accumulated in `dtype` boundary 2^31 - 1 for int32 probed at [2e9, 2e9]: int64 gives [0, 2e9, 4e9]; int32 wraps silently diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/dilated_short_conv_paged.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/dilated_short_conv_paged.py index 6ad630e90f1f..aec821713912 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/dilated_short_conv_paged.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/dilated_short_conv_paged.py @@ -69,7 +69,8 @@ class DilatedShortConvPaged(OpWrapper): sequence (`cache_indices[s] == pad_slot_id`) publishes nothing and still gets output columns -- its tokens against a zero carry, defined and meaningless. Two sequences naming one slot: the last live one's state - survives (defined, not a race). + survives (defined by construction, not a race; uncertified, as the + scheduler never names one slot twice). Fusion boundary: state read, zero fill for a fresh sequence, the dilated depthwise conv, bias, SiLU and the state publish are inside. The @@ -87,132 +88,44 @@ class DilatedShortConvPaged(OpWrapper): # 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)), + CELLS: tuple[Cell, ...] = 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, ) - + 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, - ), - *( - _ckpt( - f"MTP verify p={p} at 461 slots: {n} rows of {p} tokens", - (p,) * n, - pool_slots=ENGINE_SLOTS, - slots=_scattered(n), - x_form="token_major", - qsl_dtype=torch.int64, - ) - for p, n in ((5, 8), (6, 32), (7, 8), (8, 1), (8, 32)) - ), - ) - + 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 - ), - ), + for why, lengths in ( + ("37+1+128", (37, 1, 128)), + ("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, + ), + _ckpt( + "MTP verify p=8 at 461 slots: 32 rows of 8 tokens", + (8,) * 32, + pool_slots=ENGINE_SLOTS, + slots=_scattered(32), + x_form="token_major", + qsl_dtype=torch.int64, + ), + 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), + ), ) note = """ @@ -223,17 +136,18 @@ class DilatedShortConvPaged(OpWrapper): `[T, 10240]` normed activation, weight `[10240, 1, 4]` bf16, no bias, dilation 3 (state 9), SiLU on, the whole `[slots, 10240, 9]` pool, int32 slots, `query_start_loc_long` (int64, rebuilt from `seq_lens` under MTP), - carry flags covering the whole batch, `pad_slot_id` -1. Certified at 8, 32 - and 461 slots; at 461 bitwise identical to 8 for the same slot content. MTP - verify batches are certified at 3 and at 5..8 tokens per request, up to 32 - requests. + carry flags covering the whole batch, `pad_slot_id` -1. Certified in that + form at the 461-slot pool (prefill mixes up to 2047 tokens, 32 one-token + decode rows, MTP verify batches of 3 and of 8 tokens per request, up to 32 + requests), plus a pad row on a 6-slot pool. Preconditions: `query_start_loc[-1] == x.shape[1]`; a short one indexes past the internal history and fires a device-side assert that kills the CUDA context (loud, not guarded). An out-of-range slot reaches the same bounds check (reasoned, not probed). Pool, weight and bias share `x`'s dtype (each mismatch raises); `weight` is 3-D depthwise (2-D raises). - Kernel sizes 2-5 and dilations 1-3 certified; there is no width dispatch. + Only the checkpoint's kernel size 4 and dilation 3 are certified; there is + no width dispatch. CUDA-graph capture is certified: no host sync, all batch reads on device, a captured graph stays correct when the batch tensors' values change. @@ -242,8 +156,7 @@ class DilatedShortConvPaged(OpWrapper): Not measured: pools above 461 slots (a bf16 recurrent state gives the engine larger ones; a slot is addressed by its index alone), token counts above 8192, - dilation above 3, kernel above 5, `torch.compile`, a pool dtype unlike the - model's. + `torch.compile`, a pool dtype unlike the model's. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/extract_transpose_prefill_slice.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/extract_transpose_prefill_slice.py index 938b668b5d7b..eebc6523d738 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/extract_transpose_prefill_slice.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/extract_transpose_prefill_slice.py @@ -42,112 +42,44 @@ class ExtractTransposePrefillSlice(OpWrapper): # 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="GDN prefill, whole 10240-wide row, 1 token", + spec=dict( + num_rows=1, + num_cols=GDN_MIXED_QKV_WIDTH, + n=1, + 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=f"GDN prefill, the q|k|v column view of the in_proj, {n} tokens", - spec=dict( - num_rows=n, - num_cols=GDN_IN_PROJ_WIDTH, - col_slice=(0, GDN_MIXED_QKV_WIDTH), - n=n, - start_col=0, - width=GDN_MIXED_QKV_WIDTH, - dtype=torch.bfloat16, - ), - ) - for n in (33, 8192) + ), + 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="column-slice view: row stride wider than the extracted width", + why=f"GDN prefill, the q|k|v column view of the in_proj, {n} tokens", spec=dict( - num_rows=200, - num_cols=16384, - col_slice=(1024, 1024 + GDN_MIXED_QKV_WIDTH), - n=150, + num_rows=n, + num_cols=GDN_IN_PROJ_WIDTH, + col_slice=(0, GDN_MIXED_QKV_WIDTH), + n=n, start_col=0, width=GDN_MIXED_QKV_WIDTH, dtype=torch.bfloat16, ), - ), - ) - + tuple( - Cell( - why=f"dtype {str(dtype).removeprefix('torch.')}: elements are moved, never interpreted", - spec=dict(num_rows=64, num_cols=256, n=48, start_col=32, width=128, dtype=dtype), - ) - for dtype in ( - torch.float16, - torch.float32, - torch.float64, - torch.float8_e4m3fn, - torch.int32, - torch.int8, - torch.bool, ) - ) - + ( - Cell( - why="zero prefill tokens: an empty [width, 0] result", - spec=dict( - num_rows=64, num_cols=256, n=0, start_col=0, width=256, dtype=torch.bfloat16 - ), - ), - Cell( - why="zero width: an empty [0, n] result", - spec=dict( - num_rows=64, num_cols=256, n=1, start_col=0, width=0, dtype=torch.bfloat16 - ), - ), - ) + for n in (33, 8192) + ), ) note = """ @@ -169,6 +101,9 @@ class ExtractTransposePrefillSlice(OpWrapper): a guard). The source offset is widened to int64 in-kernel: a source offset of 2.62e9 was bitwise exact. At width 10240 the destination boundary is n > 209715, within reach of a large `max_num_tokens`. + + Not certified: dtypes other than bf16, `start_col` other than 0, a row crop + shorter than `src`, empty extents. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/flashinfer_chunk_gated_delta_rule.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/flashinfer_chunk_gated_delta_rule.py index 8a64fa713223..ec5d273c92a1 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/flashinfer_chunk_gated_delta_rule.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/flashinfer_chunk_gated_delta_rule.py @@ -73,9 +73,8 @@ class FlashinferChunkGatedDeltaRule(OpWrapper): ARCHS = frozenset({Arch.SM_103}) # spec: geometry (H, HV, K, V); seqlens; `slots`-row pool with `indices` - # (or _engine_slots); inplace (True); output_final_state; output (None | - # "buffer"); l2norm (False) with normalize_qk (True: caller-normalized). - # Every cell but the last is qwen3.8-flash-next-nvfp4's geometry. + # (or _engine_slots); output (None | "buffer"). Every cell is + # qwen3.8-flash-next-nvfp4's geometry, in place over an fp32 pool. CELLS: tuple[Cell, ...] = tuple( Cell( why=f"qwen3.8 prefill, 461-slot pool, seqlens {list(s)}", @@ -86,7 +85,7 @@ class FlashinferChunkGatedDeltaRule(OpWrapper): indices=_engine_slots(len(s)), ), ) - for s in ((1,), (64,), (200, 37, 91), (1024,), (8192,), (2048, 2048, 2048, 2048)) + for s in ((200, 37, 91), (8192,), (2048, 2048, 2048, 2048)) ) + ( Cell( why="qwen3.8 context chunk + 1-token rows [8, 1, 1, 1], 461-slot pool", @@ -107,32 +106,6 @@ class FlashinferChunkGatedDeltaRule(OpWrapper): output="buffer", ), ), - Cell( - why="qwen3.8 no in-place: final states returned, pool untouched", - spec=dict( - geometry=QWEN38_GEOMETRY, - seqlens=(96, 48), - slots=5, - indices=(3, 0), - inplace=False, - output_final_state=True, - ), - ), - Cell( - why="qwen3.8 use_qk_l2norm_in_kernel on raw q/k", - spec=dict( - geometry=QWEN38_GEOMETRY, - seqlens=(100,), - slots=3, - indices=(2,), - l2norm=True, - normalize_qk=False, - ), - ), - Cell( - why="HV/H = 2 (16 q/k heads, 32 v heads), seqlens [200, 37]", - spec=dict(geometry=(16, 32, 128, 128), seqlens=(200, 37), slots=4, indices=(3, 0)), - ), ) note = """ @@ -173,7 +146,9 @@ class FlashinferChunkGatedDeltaRule(OpWrapper): caller or the flag). Not certified: `state_workspace=`, bf16/fp16 pool, `scale` other than the - default, fp16 operands, K or V other than 128. + default, fp16 operands, K or V other than 128, other head ratios, and what the + target never passes: `inplace_indexed_state_update=False` (final states + returned) and `use_qk_l2norm_in_kernel`. Hard limits none probed. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_gating.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_gating.py index e81f6f95b7b8..d28ef96b8432 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_gating.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_gating.py @@ -51,90 +51,28 @@ class FusedGdnGating(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec: rows (or num_decodes x draft); layout "contiguous" | "target_rows" - # (the decode rows of a contiguous [T, 48] `a` projection, past T - rows - # prefill rows) | "ba_column_half" (the stock mixer's [:, 48:] of [T, 96]) - # | "ba_row_and_column" (a row slice of that); a_dtype, param_dtype, - # num_heads, beta, threshold, a_scale, fill (constant a), overlong. - 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)), + # spec: rows; layout "contiguous" | "target_rows" (the decode rows of a + # contiguous [T, 48] `a` projection, past T - rows prefill rows); a_dtype, + # a_scale, fill (constant a). + 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), (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="softplus beta 1.0, threshold 20", + spec=dict(rows=96, beta=1.0, threshold=20.0, a_scale=4.0), + ), + Cell( + why="a = 200: linear branch; exp(x) is evaluated, inf never leaks", + spec=dict(rows=4, fill=200.0, a_dtype=torch.float32), + ), ) note = """ @@ -149,9 +87,9 @@ class FusedGdnGating(OpWrapper): `a.shape`), CPU tensors (`ValueError: Pointer argument cannot be accessed from Triton`). - `a`'s row stride is free (consumed in place, bitwise equal to a packed - copy); everything else is addressed without strides. `is_valid` guards the - three violations measured to return a finite wrong `g`: a stride-2 A_log + `a`'s row stride is honoured (a row is read at `row * a.stride(0)`; the + target passes contiguous rows); everything else is addressed without + strides. `is_valid` guards the three violations measured to return a finite wrong `g`: a stride-2 A_log view moved g by 6.031e+01 and a stride-2 dt_bias by 3.832e-01; 40-element A_log / dt_bias against 48 heads moved it by 1.009e+02 / 1.010e+02 (the per-head load is masked against `a.shape[1]` only); an `a` with head diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_post_conv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_post_conv.py index c83653169fe6..83128583a121 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_post_conv.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_post_conv.py @@ -63,66 +63,30 @@ class FusedGdnPostConv(OpWrapper): ARCHS = frozenset({Arch.SM_103}) # spec: num_prefill, num_decode; layout "contiguous" | "target_mixed" - # (decode a row slice of a [T, 10240] projection, as the target's split - # makes it) | "production" (the stock mixer's views: decode a row slice of - # a column slice of [T, 10240 + 6144], a/b column halves of [T, 96]) | - # "prefill_view" (prefill a `.t()` of a wider buffer, stride (1, 16384)) | + # (decode a row slice of a [T, 10240] projection, a/b contiguous [T, 48]) | # "in_proj" (the target's one [T, 16480] input projection: decode the - # q|k|v column view of its decode rows, a/b its column views); - # beta_dtype, gate_dtype (A_log/dt_bias), dt_bias_fill. + # q|k|v column view of its decode rows, a/b its column views); 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), + why=f"target in_proj views, {p} prefill + {d} decode tokens", + spec=dict(num_prefill=p, num_decode=d, layout="in_proj"), ) - 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=f"target in_proj views, {p} prefill + {d} decode tokens", - spec=dict(num_prefill=p, num_decode=d, layout="in_proj"), - ) - for p, d in ((200, 32), (2048, 0), (0, 128)) - ), - Cell( - why="stock mixer views: strided decode, a/b column halves", - spec=dict(num_prefill=96, num_decode=12, layout="production"), - ), - Cell( - why="prefill a non-contiguous .t() view, read through both strides", - spec=dict(num_prefill=96, num_decode=0, layout="prefill_view"), - ), - Cell( - why="beta_dtype bf16 (the MTP verify-path convention)", - spec=dict(num_prefill=37, num_decode=0, beta_dtype=torch.bfloat16), - ), - Cell( - why="A_log/dt_bias bf16, as stored in the checkpoint", - spec=dict(num_prefill=64, num_decode=8, gate_dtype=torch.bfloat16), - ), - Cell( - why="every head below softplus_threshold", - spec=dict(num_prefill=64, num_decode=0, dt_bias_fill=SOFTPLUS_THRESHOLD - 8.0), - ), - Cell( - why="every head above softplus_threshold: the linear branch", - spec=dict(num_prefill=64, num_decode=0, dt_bias_fill=SOFTPLUS_THRESHOLD + 8.0), - ), - Cell( - why="zero tokens: five empty tensors, no launch", - spec=dict(num_prefill=0, num_decode=0), - ), - ) + for p, d in ((200, 32), (2048, 0)) + ), + Cell( + why="target mixed call: decode a row slice of the projection", + spec=dict(num_prefill=200, num_decode=32, layout="target_mixed"), + ), + Cell(why="prefill only, 1 token", spec=dict(num_prefill=1, num_decode=0)), + Cell( + why="mixed batch, 1 prefill + 1 decode tokens", + spec=dict(num_prefill=1, num_decode=1), + ), + Cell( + why="every head above softplus_threshold: the linear branch", + spec=dict(num_prefill=64, num_decode=0, dt_bias_fill=SOFTPLUS_THRESHOLD + 8.0), + ), ) note = """ @@ -136,11 +100,11 @@ class FusedGdnPostConv(OpWrapper): `extract_transpose_prefill_slice`'s contiguous output; `decode` and `a`/`b` column views (row stride 16480) of the layer's one `[T, 16480]` input-projection output. The stock mixer's verify call passes - `beta_dtype=b.dtype`. + `beta_dtype=b.dtype` (not certified here). Preconditions: `prefill.shape[0]` and `decode.shape[1]` >= qkv_dim (extra trailing features ignored; a short one is read out of bounds); `decode` - has `prefill`'s dtype. `a`/`b` may have extra columns (the production + has `prefill`'s dtype. `a`/`b` may have extra columns (the in_proj case) and are read through their own strides. `head_k_dim`/`head_v_dim` are rounded up to a power of two and masked. @@ -158,7 +122,9 @@ class FusedGdnPostConv(OpWrapper): (`tl.where` is a select); harmless below fp32 overflow at x ~ 88, which no checkpoint `a + dt_bias` reaches. - Not certified: fp16/fp32 activations, other head geometries. + Not certified: fp16/fp32 activations, other head geometries, a bf16 + `beta_dtype`, bf16 `A_log`/`dt_bias`, a non-contiguous `prefill`, an empty + batch. Hard limits none probed. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_cached_replay_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_cached_replay_update.py index d131919c8078..6edf1d4e38e8 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_cached_replay_update.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_cached_replay_update.py @@ -122,55 +122,17 @@ class FusedRecurrentGatedDeltaRuleCachedReplayUpdate(OpWrapper): # dup (rows given the SAME slot, as CUDA-graph padding repeats one dummy slot; # only the other rows and slots are held to the reference). CELLS: tuple[Cell, ...] = ( + # In-kernel commit (below 16 requests), one cell per token tile BT = next_pow2(T). Cell(why="small n=1 T=2: one sequence, MTP-1", spec=_cell_spec("small", 1, 2, steps=12)), - Cell(why="small n=4 T=3", spec=_cell_spec("small", 4, 3)), - Cell( - why="small n=8 T=4 L=2 S=461: slots 0 and 460 of a 461-slot axis", - spec=_cell_spec("small", 8, 4, slots=ENGINE_SLOTS, edge_slots=True), - ), - Cell( - why="small n=15 T=2: largest batch before the split", - spec=_cell_spec("small", 15, 2, steps=12), - ), - Cell( - why="small n=8 T=4 fresh: first verify after prefill", - spec=_cell_spec("small", 8, 4, start="fresh", steps=10), - ), - Cell( - why="small n=6 T=3 pad: rows 1, 4 at slot -1, prefix output", - spec=_cell_spec("small", 6, 3, pad=(1, 4), output="prefix"), - ), Cell( why="small n=8 T=3 dup: rows 5..7 share one padding slot", spec=_cell_spec("small", 8, 3, dup=(5, 6, 7), steps=6), ), Cell( - why="small n=4 T=2 unpacked: q/k/v views alone", - spec=_cell_spec("small", 4, 2, packed=False, steps=12), - ), - Cell( - why="small n=4 T=4 PDL on: launch_with_pdl=True, the kernel's other compile", - spec=_cell_spec("small", 4, 4, pdl=True), - ), - Cell( - why="small n=5 T=3 strided a/b: halves of [n*t, 96]", - spec=_cell_spec("small", 5, 3, ab="strided"), - ), - Cell( - why="all_layer n=16 T=4: the partition threshold", spec=_cell_spec("all_layer", 16, 4) - ), - Cell( - why="all_layer n=24 T=3 L=2 S=461: slots 0 and 460 of a 461-slot axis", - spec=_cell_spec("all_layer", 24, 3, slots=ENGINE_SLOTS, edge_slots=True), - ), - Cell( - why="all_layer n=32 T=4: the envelope's 32 requests", - spec=_cell_spec("all_layer", 32, 4), - ), - Cell( - why="all_layer n=32 T=3 L=36 S=40: every GDN layer, one commit", - spec=_cell_spec("all_layer", 32, 3, layers=NUM_GDN_LAYERS, slots=40, steps=6), + why="small n=4 T=8 full accept: pnat+T==16 appends, then wraps", + spec=_cell_spec("small", 4, 8, start="fresh", accept="full", steps=8), ), + # Partitioned call plus the all-layer commit (16 requests and up). Cell( why="all_layer n=16 T=3 L=8 S=461: slot 460's offset is past 2^31 elements", spec=_cell_spec( @@ -183,95 +145,11 @@ class FusedRecurrentGatedDeltaRuleCachedReplayUpdate(OpWrapper): steps=3, ), ), - Cell( - why="all_layer n=20 T=4 fresh: after prefill", - spec=_cell_spec("all_layer", 20, 4, start="fresh", steps=10), - ), - Cell( - why="all_layer n=18 T=3 pad: rows 0, 17 at slot -1", - spec=_cell_spec("all_layer", 18, 3, pad=(0, 17), output="prefix"), - ), - Cell( - why="all_layer n=16 T=3 mixed-batch form: views, offset out", - spec=_cell_spec("all_layer", 16, 3, packed=False, output="suffix"), - ), - # ---- MTP draft lengths 4..7: T = 5..8, the kernel's ceiling ---- - Cell(why="small n=1 T=5: one sequence, MTP-4", spec=_cell_spec("small", 1, 5, steps=10)), - Cell(why="small n=8 T=6: MTP-5", spec=_cell_spec("small", 8, 6)), - Cell( - why="small n=15 T=7: MTP-6, largest batch before the split", - spec=_cell_spec("small", 15, 7), - ), - Cell(why="small n=1 T=8: one sequence, MTP-7", spec=_cell_spec("small", 1, 8, steps=10)), - Cell(why="small n=8 T=8: MTP-7", spec=_cell_spec("small", 8, 8)), - Cell( - why="small n=15 T=8: MTP-7, largest batch before the split", - spec=_cell_spec("small", 15, 8), - ), - Cell( - why="small n=8 T=8 S=461: slots 0 and 460 of a 461-slot axis", - spec=_cell_spec("small", 8, 8, slots=ENGINE_SLOTS, edge_slots=True), - ), - Cell( - why="small n=4 T=8 full accept: pnat+T==16 appends, then wraps", - spec=_cell_spec("small", 4, 8, start="fresh", accept="full", steps=8), - ), - Cell( - why="small n=8 T=8 one accepted: a=1 every step", - spec=_cell_spec("small", 8, 8, accept="one", steps=12), - ), - Cell( - why="small n=6 T=8 mixed-batch form: views, offset out, pad", - spec=_cell_spec("small", 6, 8, packed=False, output="suffix", pad=(2,)), - ), - Cell(why="all_layer n=16 T=5: MTP-4 at the threshold", spec=_cell_spec("all_layer", 16, 5)), - Cell( - why="all_layer n=24 T=6 S=461: MTP-5, slots 0 and 460", - spec=_cell_spec("all_layer", 24, 6, slots=ENGINE_SLOTS, edge_slots=True), - ), - Cell(why="all_layer n=32 T=7: MTP-6 at 32 requests", spec=_cell_spec("all_layer", 32, 7)), - Cell(why="all_layer n=16 T=8: MTP-7 at the threshold", spec=_cell_spec("all_layer", 16, 8)), - Cell(why="all_layer n=32 T=8: MTP-7, 256 rows", spec=_cell_spec("all_layer", 32, 8)), - Cell( - why="all_layer n=32 T=8 L=36 S=40: every GDN layer, one commit", - spec=_cell_spec("all_layer", 32, 8, layers=NUM_GDN_LAYERS, slots=40, steps=6), - ), Cell( why="all_layer n=32 T=8 full accept: commit every second step", spec=_cell_spec("all_layer", 32, 8, start="fresh", accept="full", steps=8), ), - Cell( - why="all_layer n=16 T=8 mixed-batch form: views, offset out", - spec=_cell_spec("all_layer", 16, 8, packed=False, output="suffix"), - ), - Cell( - why="all_layer n=20 T=8 fresh pad: rows 0, 19 at slot -1", - spec=_cell_spec("all_layer", 20, 8, start="fresh", pad=(0, 19), output="prefix"), - ), # ---- a bf16 checkpoint pool (`mamba_ssm_cache_dtype=bfloat16`) ---- - _bf16_cell("small n=1 T=2: one sequence, MTP-1", "small", 1, 2, steps=12), - _bf16_cell("small n=8 T=4", "small", 8, 4), - _bf16_cell("small n=15 T=3: largest batch before the split", "small", 15, 3), - _bf16_cell("small n=8 T=4 fresh: first verify after prefill", "small", 8, 4, start="fresh"), - _bf16_cell( - "small n=6 T=3 pad: rows 1, 4 at slot -1, prefix output", - "small", - 6, - 3, - pad=(1, 4), - output="prefix", - ), - _bf16_cell( - "small n=8 T=4 S=461: slots 0 and 460 of a 461-slot axis", - "small", - 8, - 4, - slots=ENGINE_SLOTS, - edge_slots=True, - ), - _bf16_cell("small n=1 T=5: MTP-4", "small", 1, 5, steps=10), - _bf16_cell("small n=1 T=8: MTP-7, one sequence", "small", 1, 8, steps=10), - _bf16_cell("small n=8 T=8: MTP-7", "small", 8, 8), _bf16_cell( "small n=4 T=8 full accept: pnat+T==16 appends, then wraps", "small", @@ -281,9 +159,6 @@ class FusedRecurrentGatedDeltaRuleCachedReplayUpdate(OpWrapper): accept="full", steps=8, ), - _bf16_cell("all_layer n=16 T=4: the partition threshold", "all_layer", 16, 4), - _bf16_cell("all_layer n=24 T=6: MTP-5", "all_layer", 24, 6), - _bf16_cell("all_layer n=32 T=8: MTP-7, 256 rows", "all_layer", 32, 8), _bf16_cell( "all_layer n=32 T=3 L=36 S=40: every GDN layer, one commit", "all_layer", @@ -293,23 +168,6 @@ class FusedRecurrentGatedDeltaRuleCachedReplayUpdate(OpWrapper): slots=40, steps=6, ), - _bf16_cell( - "all_layer n=32 T=8 full accept: commit every second step", - "all_layer", - 32, - 8, - start="fresh", - accept="full", - steps=8, - ), - _bf16_cell( - "all_layer n=16 T=8 mixed-batch form: views, offset out", - "all_layer", - 16, - 8, - packed=False, - output="suffix", - ), ) note = """ @@ -372,14 +230,14 @@ class FusedRecurrentGatedDeltaRuleCachedReplayUpdate(OpWrapper): Bit for bit the same call: the in-kernel commit and the partitioned call plus the all-layer commit; `packed_qkv` and the separate q/k/v (whose contents are never read when `packed_qkv` is given); `launch_with_pdl` on and off; a second - run; CUDA-graph replay on new data (the test captures every T = 2..8 at N = 1, 8, 15, 16 - and 32). + run; CUDA-graph replay on new data (the test captures N = 1 at T = 2 and N = 32 at T = 8 + over an fp32 checkpoint, N = 16 at T = 4 over a bf16 one). - T = 5..8 share one compiled token tile (`BT = next_pow2(T) = 8`). At + T = 5..8 share one compiled token tile (`BT = next_pow2(T) = 8`); the cells drive + each tile (T = 2, 3 and 8) and both checkpoint dtypes, not every pairing of the two. At T = 8 and history 16 a slot whose every draft is accepted appends rows 8..15 at `pnat + T == 16` without wrapping and wraps on the next step, so it commits every second - step; the `accept="full"` cells drive that boundary (in-kernel and all-layer commit), the - `accept="one"` cell the slowest fill. + step; the `accept="full"` cells drive that boundary (in-kernel and all-layer commit). Loud refusals: work items without `use_all_layer_commit` (ValueError) and the reverse (ValueError); one of `replay_work_items` / `n_writes` alone diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_update.py index 13209d0dd5b3..6e57e396efcf 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_update.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_update.py @@ -27,11 +27,7 @@ #: Rows of the engine's per-step buffer (`max_num_sequences`), as the #: runtime-value dumps record `intermediate_ssm_states`: [32, p, 48, 128, 128]. ENGINE_VERIFY_ROWS = 32 -#: A recurrent-pool slot count of a full-size engine run (a bf16 pool holds more slots in -#: the same memory). The target gathers a batch-scoped copy before this call, so the -#: write-back mode is bracketed here rather than driven. -ENGINE_SLOTS = 461 -#: Pool and buffer rows of the small cells. +#: Default pool and buffer rows of the test's builds. SLOTS = 8 @@ -39,10 +35,6 @@ 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)`): @@ -80,115 +72,23 @@ class FusedRecurrentGatedDeltaRuleUpdate(OpWrapper): # 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)) - ) - + tuple( - _verify( - n, - p, - f"bf16 state: target verify n={n} p={p}, n-row gather, 32-row buffer", - pool_slots=n, - cache_rows=ENGINE_VERIFY_ROWS, - beta="bf16", - qkv="contiguous", - pool_dtype=torch.bfloat16, - cache_dtype=torch.bfloat16, - ) - for n, p in ((1, 2), (8, 2), (16, 2), (32, 2), (8, 3), (8, 4)) - ) - + ( - _verify( - 4, - 2, - "bf16 state: verify, -1-padded rows naming 31 and 0", - pool_slots=ENGINE_VERIFY_ROWS, - cache_rows=ENGINE_VERIFY_ROWS, - slots=(31, -1, 0, -1), - pool_dtype=torch.bfloat16, - cache_dtype=torch.bfloat16, - ), - ) - + ( - _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) + CELLS: tuple[Cell, ...] = tuple( + _verify( + n, + p, + f"{label}target verify: n={n} p={p}, n-row gather, 32-row buffer", + pool_slots=n, + cache_rows=ENGINE_VERIFY_ROWS, + beta="bf16", + qkv="contiguous", + pool_dtype=dtype, + cache_dtype=dtype, ) - + ( - _write_back(3, 4, "write-back into a prefix of a longer output", output="prefix"), - _write_back(2, 1, "write-back P=1 into a prefix of an output", output="prefix"), - _write_back(3, 3, "write-back slot -1: no slot written", slots=(2, -1, 5)), - _write_back(2, 4, "write-back, every slot -1", slots=(-1, -1)), - _write_back(3, 4, "write-back into a bf16 pool", pool_dtype=torch.bfloat16), - _write_back(2, 1, "write-back P=1 into a bf16 pool", pool_dtype=torch.bfloat16), - _write_back(3, 1, "bf16 beta (the mixer's dtype), P=1", slots=(2, 5, 0), beta="bf16"), - _write_back(3, 4, "bf16 beta (the mixer's dtype), P=4", slots=(2, 5, 0), beta="bf16"), - _write_back(2, 3, "per-channel beta [N, P, HV, V]", slots=(1, 2), beta="channel"), - _write_back(2, 3, "fp32 output buffer against bf16 q", slots=(1, 2), output="fp32"), - _write_back(2, 3, "int64 initial_state_indices", slots=(1, 2), index_dtype=torch.int64), - _write_back( - 2, - 3, - "a [layers, S, ...] pool's middle layer", - slots=(1, 4), - pool_view="layer_major", - ), - Cell(why="both sinks at once", spec=dict(mode="both", n=2, p=4)), - _write_back(3, None, "varlen: ragged lengths 1/5/2", lengths=(1, 5, 2)), - _write_back( - 3, - None, - "varlen: uniform lengths, scattered slots", - lengths=(4, 4, 4), - slots=(2, 6, 3), - ), - _write_back(2, None, "varlen: int32 cu_seqlens", lengths=(2, 4), cu_dtype=torch.int32), + for dtype, label, shapes in ( + (torch.float32, "", ((1, 2), (32, 2), (8, 4))), + (torch.bfloat16, "bf16 state: ", ((1, 2), (32, 2), (8, 3))), ) + for n, p in shapes ) note = """ @@ -231,8 +131,13 @@ class FusedRecurrentGatedDeltaRuleUpdate(OpWrapper): 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. + Certified: the call site's form alone -- the per-step buffer sink with the pool + untouched, at 1, 8 and 32 sequences and P = 2..4, over an fp32 or a bf16 pool and buffer. + Not certified: the write-back sink (`disable_state_update=False`), alone or with the + buffer; the varlen form (`cu_seqlens`); negative, scattered or int64 indices; no pool; + fp32 or per-channel `beta`; q/k/v as views of the packed rows; an `output` buffer; + `disable_output_calculation=True`; `K` or `V` other than 128; fp16/fp32 q/k/v; P + outside 2..4; N above 32. Hard limits N * HV gridDim.z boundary 65535 probed at N=1300 (62400): runs; diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_sigmoid_gating_delta_rule_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_sigmoid_gating_delta_rule_update.py index e487455c7ebb..a3d336b7b984 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_sigmoid_gating_delta_rule_update.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_sigmoid_gating_delta_rule_update.py @@ -19,9 +19,6 @@ 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 #: A recurrent-pool slot count of a full-size engine run. The kernels address a slot by #: its index alone, so a pool of another size (a bf16 pool holds more slots in the same #: memory) takes the same path. The target passes the layer's whole pool unsliced. @@ -29,8 +26,10 @@ def _scattered(n_seqs: int) -> tuple[int, ...]: - """`n_seqs` distinct engine-pool slots, both ends included.""" - return ((0, ENGINE_SLOTS - 1) + tuple(i * 13 + 4 for i in range(n_seqs)))[:n_seqs] + """`n_seqs` distinct engine-pool slots, both ends first, then a stride-13 walk over the + interior (13 is coprime to its `ENGINE_SLOTS - 2` slots).""" + rest = tuple((i * 13 + 4) % (ENGINE_SLOTS - 2) + 1 for i in range(max(n_seqs - 2, 0))) + return ((0, ENGINE_SLOTS - 1) + rest)[:n_seqs] class FusedSigmoidGatingDeltaRuleUpdate(OpWrapper): @@ -72,179 +71,60 @@ class FusedSigmoidGatingDeltaRuleUpdate(OpWrapper): # contiguous copies | "rows": cu_seqlens None, one token per row, q/k/v # [N, 1, heads, 128] views of the packed [N, 10240] block whose token # stride is the row stride, and a/b [N, 1, 48] -- the target's bf16 call). - 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) - ) - + tuple( - Cell( - why=f"target decode call, in_proj column views, N={n}", - spec=dict(seq_lens=(1,) * n, slots=slots, indices=indices, ab="in_proj"), - ) - for n, slots, indices in ( - (1, ENGINE_SLOTS, _scattered(1)), - (32, ENGINE_SLOTS, _scattered(32)), - (128, 128, None), - ) - ) - + ( - 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"target bf16 decode call, rows form N={n}, 461-slot pool", - spec=dict( - seq_lens=(1,) * n, - slots=ENGINE_SLOTS, - indices=_scattered(n), - ab="separate", - layout="rows", - pool_dtype=torch.bfloat16, - ), - ) - for n in (1, 3, 10, 11, 16, 32) - ) - + ( - Cell( - why="bf16 rows form, pad rows (-1) between slot 0 and 460", - spec=dict( - seq_lens=(1, 1, 1, 1), - slots=ENGINE_SLOTS, - indices=(0, -1, ENGINE_SLOTS - 1, -1), - ab="separate", - layout="rows", - pool_dtype=torch.bfloat16, - ), - ), - Cell( - why="fp32 pool, rows form N=8, 461-slot pool", - spec=dict( - seq_lens=(1,) * 8, - slots=ENGINE_SLOTS, - indices=_scattered(8), - ab="separate", - layout="rows", - ), - ), - ) - + 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)), + CELLS: tuple[Cell, ...] = tuple( + Cell( + why=f"target bf16 decode call, rows form N={n}, 461-slot pool", + spec=dict( + seq_lens=(1,) * n, + slots=ENGINE_SLOTS, + indices=_scattered(n), + ab="separate", + layout="rows", + pool_dtype=torch.bfloat16, ), ) - + 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), + for n in (1, 32, 128) + ) + ( + Cell( + why="bf16 rows form, pad rows (-1) between slot 0 and 460", + spec=dict( + seq_lens=(1, 1, 1, 1), + slots=ENGINE_SLOTS, + indices=(0, -1, ENGINE_SLOTS - 1, -1), + ab="separate", + layout="rows", + pool_dtype=torch.bfloat16, ), - ) - + 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"), - ), - ) + ), + Cell( + why="fp32 pool, in_proj column views, N=128: the block test's reference", + spec=dict(seq_lens=(1,) * 128, slots=128, ab="in_proj"), + ), ) note = """ Metadata consumed: none. - Call sites (`_gdn_block` decode), one token per sequence, int32 slots, the - layer's whole pool `[slots, 48, 128, 128]` at `mamba_ssm_cache_dtype`: - - fp32 pool: `(a_log, a, dt_bias, 1.0, 20.0, q, k, v, b, ssm_states, - si_all, None, True, qsl_long, None)` -- q/k/v `[1, T, heads, 128]` - reshaped column views of the conv output (read through `stride(1)`, no - repack), and `a`/`b`, all column views of the layer's one `[T, 16480]` - input-projection output (row stride 16480; the conv wrote q|k|v in - place), int64 `cu_seqlens`; - - bf16 pool: the rows form, `cu_seqlens` None -- q/k/v `[T, 1, heads, - 128]` split from the whole `[T, 10240]` block reshaped to `[T, 1, 80, - 128]`, so the token stride is the row stride (`is_valid` refuses other - strides), `a`/`b` `[T, 1, 48]`. - Certified at 24 and 461 slots; at 461 output and state are bitwise - identical to 24 for the same slot content. The rows form is bitwise the - packed Triton call at either pool dtype and runs inside a CUDA graph. + Call site (`_gdn_block` decode over a bf16 pool), one token per sequence, + int32 slots, the layer's whole pool `[slots, 48, 128, 128]`: the rows form, + `cu_seqlens` None -- q/k/v `[T, 1, heads, 128]` split from the whole + `[T, 10240]` block reshaped to `[T, 1, 80, 128]`, so the token stride is the + row stride (`is_valid` refuses other strides), `a`/`b` `[T, 1, 48]`. The + cells drive it at 1, 32 and 128 rows of a 461-slot pool, and with CUDA-graph + pad rows (-1). One more cell certifies the packed form the GDN block test + uses as its fp32-pool reference: `(a_log, a, dt_bias, 1.0, 20.0, q, k, v, + b, ssm_states, si_all, None, True, qsl_long, None)`, q/k/v `[1, T, heads, + 128]` and `a`/`b` column views of one `[T, 16480]` input-projection output + (row stride 16480), int64 `cu_seqlens`, at 128 rows. The rows form is + bitwise the packed Triton call at either pool dtype and runs inside a CUDA + graph. Branch: the callable takes a FlashInfer CuTe-DSL kernel (`gdn_decode_bf16_state.gated_delta_rule`) when the pool is bf16, `cu_seqlens` is given and `T == N` (head dims 128, supported arch); everything else runs the vendored Triton kernel, and only the Triton - branch is certified. The target's two forms both stay on it, which the - test proves with the FlashInfer arm armed to raise. The bf16 call takes + branch is certified. The target's rows form and the reference's fp32 packed form both + stay on it, which the test proves with the FlashInfer arm armed to raise. The bf16 call takes the rows form because FlashInfer is not batch-invariant: at 48 value heads it runs one kernel up to N = 10 and another from N = 11, and one fixed row's single-step output and state differ between the two, so a @@ -276,7 +156,9 @@ class FusedSigmoidGatingDeltaRuleUpdate(OpWrapper): Not certified: the FlashInfer branch, K or V other than 128, fp16/fp32 operands, an fp16 pool, `softplus_beta`/`softplus_threshold` other than - 1.0/20.0, `B > 1` with `cu_seqlens`. + 1.0/20.0, `B > 1` with `cu_seqlens`; the rows form above 32 rows; the packed + form over a bf16 pool, with more than one token per sequence, with an + `output` buffer, with `l2norm` off, without a pool or on a layer view of one. Hard limits none probed. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_in_proj_conv1d_pdl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_in_proj_conv1d_pdl.py index b31ddad3daab..926c3a9589f8 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_in_proj_conv1d_pdl.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_in_proj_conv1d_pdl.py @@ -321,7 +321,6 @@ class GdnInProjConv1dPdl(OpWrapper): # spec: rows, slots (tuple per row; -1 is the pad slot), pdl, pool (slots in the pool). CELLS: tuple[Cell, ...] = ( _cell("decode 1 row, slot 37 of the 461-slot pool", 1, (37,)), - _cell("decode 1 row, the last slot of the pool", 1, (460,)), _cell("decode 2 rows, two slots", 2, (5, 211)), _cell("decode 2 rows, one a pad row (-1)", 2, (17, -1)), _cell("decode 1 row, PDL off", 1, (99,), pdl=False), diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_prefill_conv_post.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_prefill_conv_post.py index 29c567150675..670e126997de 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_prefill_conv_post.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_prefill_conv_post.py @@ -774,18 +774,11 @@ class GdnPrefillConvPost(OpWrapper): # spec: lengths; slots (pool rows used, None -> _engine_slots); pool_slots; # initial; layout "in_proj" (x and a/b column views of one [T, 16480] - # buffer, the target's call) | "contiguous" (x [T, 10240], a/b [T, 48]) | - # "rows" (x/a/b the first rows of a longer projection, as a mixed step's - # context half); pad_slot_id; dt_bias_fill. + # buffer, the target's call) | "rows" (x/a/b the first rows of a longer + # projection, as a mixed step's context half); pad_slot_id; dt_bias_fill. CELLS: tuple[Cell, ...] = ( Cell( - why="target call: 37+1+128 tokens, in_proj views, 8-slot pool", - spec=dict( - lengths=(37, 1, 128), pool_slots=8, slots=(3, 0, 5), initial=(True, False, True) - ), - ), - Cell( - why="engine 461-slot pool: slots 460, 0, 230", + why="target call, engine 461-slot pool: 37+1+128 tokens at slots 460, 0, 230", spec=dict( lengths=(37, 1, 128), pool_slots=ENGINE_SLOTS, @@ -793,10 +786,6 @@ class GdnPrefillConvPost(OpWrapper): initial=(True, False, True), ), ), - Cell( - why="one 8192-token context chunk with carry", - spec=dict(lengths=(8192,), pool_slots=ENGINE_SLOTS, initial=(True,)), - ), Cell( why="max_num_tokens 16384 over 3 contexts (the con=128 mixed step)", spec=dict( @@ -814,54 +803,19 @@ class GdnPrefillConvPost(OpWrapper): ), ), Cell( - why="128 short contexts (the batch ceiling) over the 461-slot pool", + why="128 short contexts (the batch ceiling), some shorter than the window", spec=dict( lengths=tuple(1 + (i * 7) % 64 for i in range(128)), pool_slots=ENGINE_SLOTS, initial=tuple(i % 3 != 0 for i in range(128)), ), ), - Cell( - why="contexts shorter than the window: carry part shifted, part tokens", - spec=dict( - lengths=(2, 1, 2, 3), - pool_slots=6, - slots=(0, 1, 2, 5), - initial=(True, True, False, True), - ), - ), - Cell( - why="a zero-token context republishes its carry unchanged", - spec=dict( - lengths=(0, 40, 0), pool_slots=4, slots=(2, 3, 1), initial=(True, False, False) - ), - ), - Cell( - why="contiguous x [T, 10240] and a/b [T, 48]", - spec=dict( - lengths=(64, 9), - pool_slots=4, - slots=(1, 2), - initial=(True, True), - layout="contiguous", - ), - ), Cell( why="pad_slot_id -1: that context neither reads nor writes the pool", spec=dict( lengths=(16, 8, 16), pool_slots=4, slots=(0, -1, 2), initial=(True, True, True) ), ), - Cell( - why="every head below softplus_threshold", - spec=dict( - lengths=(64,), - pool_slots=2, - slots=(1,), - initial=(True,), - dt_bias_fill=SOFTPLUS_THRESHOLD - 8.0, - ), - ), Cell( why="every head above softplus_threshold: the linear branch", spec=dict( @@ -914,7 +868,7 @@ class GdnPrefillConvPost(OpWrapper): `1 + e <= 2^126`; below acc = -87.3 it returns -0 where those return a value under 1e-36. q/k norm: the 128-channel fp32 sum of squares as each thread's pairs plus a 32-lane butterfly, `sqrt.approx.ftz`, `div.full`. - Against the unfused path (entry test, every cell): + Against the unfused path (entry test, at the target-call cell): q/k/v within one bf16 ulp on 1e-6..1e-4 of elements, g/beta within the fp32 band, conv pool bitwise. @@ -936,7 +890,9 @@ class GdnPrefillConvPost(OpWrapper): Not checked: `cache_indices` within the pool and distinct; `query_start_loc` non-decreasing and ending at `T`. - Not certified: fp16/fp32 activations, other head geometries, bias. + Not certified: fp16/fp32 activations, other head geometries, bias, a row + stride of `x` other than the projection's 16480 (the conv pass compiles one + kernel per row stride), zero-token contexts. Hard limits element offsets 64-bit for the rows of x and the outputs; weight and diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/pack_gdn_decode_qkv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/pack_gdn_decode_qkv.py index 175b7f695dfd..5ff86c1442bf 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/pack_gdn_decode_qkv.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/pack_gdn_decode_qkv.py @@ -50,74 +50,15 @@ class PackGdnDecodeQkv(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec: num_tokens, or (num_decodes, draft) for a verify batch; layout - # "contiguous" | "conv_dim_major" (the conv's [B, C, P] output transposed - # and flattened, as the stock mixer builds it; a view for B == 1) | - # "row_slice_of_projection" | "column_stride_2" | "transposed_view" | - # "wide" (trailing columns past qkv_dim); dtype; heads. - CELLS: tuple[Cell, ...] = ( - tuple( - Cell(why=f"{n} tokens, contiguous{note}", spec=dict(num_tokens=n)) - for n, note in ( - (1, ""), - (2, ""), - (5, ""), - (15, ": one short of the 16-token block"), - (16, ": exactly one 16-token block"), - (17, ": one past the 16-token block"), - (32, ": two full blocks"), - (127, ": a partial final block"), - ) + # spec: num_tokens, or (num_decodes, draft) for a verify batch, whose rows + # are the token-major `[n * p, qkv_dim]` view of the conv's `[n, C, p]` output. + CELLS: tuple[Cell, ...] = tuple( + Cell( + why=f"target verify call: {b} seqs x {p} tokens, token-major rows", + spec=dict(num_decodes=b, draft=p), ) - + tuple( - Cell( - why=f"target verify call: {b} seqs x {p} tokens, token-major rows", - spec=dict(num_decodes=b, draft=p), - ) - for b, p in ((1, 2), (8, 2), (32, 2), (8, 4)) - ) - + tuple( - Cell( - why=f"stock mixer verify: {b} x {p} off the [B, C, P] conv output", - spec=dict(num_decodes=b, draft=p, layout="conv_dim_major"), - ) - for b, p in ((1, 1), (4, 1), (1, 2), (5, 2), (3, 3), (8, 4), (16, 2), (2, 5)) - ) - + ( - Cell( - why="row slice of a [T, qkv_dim + z] projection, stride (16384, 1)", - spec=dict(num_tokens=24, layout="row_slice_of_projection"), - ), - Cell( - why="column-stride-2 view, stride (20480, 2)", - spec=dict(num_tokens=24, layout="column_stride_2"), - ), - Cell( - why="transposed dim-major view, stride (1, 24)", - spec=dict(num_tokens=24, layout="transposed_view"), - ), - Cell( - why="6144 trailing columns past qkv_dim are dropped", - spec=dict(num_tokens=24, layout="wide"), - ), - ) - + tuple( - Cell( - why=f"dtype {str(dtype).removeprefix('torch.')} is carried, not chosen", - spec=dict(num_tokens=11, dtype=dtype), - ) - for dtype in (torch.bfloat16, torch.float16, torch.float32) - ) - + tuple( - Cell(why=f"geometry {heads}{note}", spec=dict(num_tokens=13, heads=heads)) - for heads, note in ( - ((16, 128, 32, 128), ""), - ((4, 64, 8, 64), ""), - ((2, 96, 6, 80), ": head dims off a power of two"), - ) - ) - + (Cell(why="zero tokens: three empty tensors, no launch", spec=dict(num_tokens=0)),) - ) + for b, p in ((1, 2), (8, 4), (32, 2)) + ) + (Cell(why="127 tokens, contiguous: a partial final block", spec=dict(num_tokens=127)),) note = """ Metadata consumed: none. Stateless; sequence boundaries and draft grouping diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/reset_recurrent_state_rows.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/reset_recurrent_state_rows.py index 7b4467a2cf7c..086581080989 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/reset_recurrent_state_rows.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/reset_recurrent_state_rows.py @@ -59,89 +59,24 @@ class ResetRecurrentStateRows(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec: slots (pool rows), indices, flags, index_dtype (int32), flag_dtype - # (bool), rec_dtype (fp32), conv (True), rec_shape/conv_shape (checkpoint), - # layout: "plain" | "row_padding" (stride(0) > row) | "layer_view" - # (pool[1] of a [3, slots, ...] allocation); `indices="perm"` is a random - # permutation of every slot with random flags. + # spec: slots (pool rows), indices, flags (bool), index_dtype (int32), + # rec_dtype (fp32), conv (True), layout: "plain" | "row_padding" + # (stride(0) > row) | "layer_view" (pool[1] of a [3, slots, ...] allocation). 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="bf16 recurrent pool 461, 4 requests, slot 0 reset", - spec=dict(_engine_case(4, flip=False), rec_dtype=torch.bfloat16), - ), Cell( why="bf16 recurrent pool 461, 32 requests, slot 0 reset", spec=dict(_engine_case(32, flip=False), rec_dtype=torch.bfloat16), ), Cell( - why="bf16 recurrent pool 461, 32 requests, flags flipped: slot 460 reset", - spec=dict(_engine_case(32, flip=True), rec_dtype=torch.bfloat16), - ), - 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", + why="negative and out-of-range slots skipped, int32", 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, + slots=6, indices=(-1, 6, 23, -99999, 2), flags=(False,) * 5, index_dtype=torch.int32 ), - ) - 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"), @@ -150,36 +85,14 @@ class ResetRecurrentStateRows(OpWrapper): why="pool[layer] of a [layers, slots, ...] allocation", spec=dict(slots=4, indices=(2, 0), flags=(False, True), layout="layer_view"), ), - Cell( - why="empty batch launches nothing", - spec=dict(slots=6, indices=(), flags=()), - ), - ) - CELLS += tuple( - Cell( - why=f"recurrent pool {str(d)[6:]} (--mamba_ssm_cache_dtype)", - spec=dict(slots=3, indices=(1,), flags=(False,), rec_dtype=d), - ) - for d in (torch.bfloat16, torch.float16) - ) - CELLS += ( - Cell( - why="conv row wider than the recurrent row", - spec=dict(slots=3, indices=(1,), flags=(False,), rec_shape=(8,), conv_shape=(5000,)), - ), - Cell( - why="64-slot random permutation: slot-indexed, not position-indexed", - spec=dict(slots=64, indices="perm", flags="random"), - ), ) note = """ Metadata consumed: none. At the call site `state_indices` / `has_initial_states` are the context prefix of `Mamba2Metadata`'s int32 / bool buffers, and the pools are the layer's views of the cache manager's - `[layers, slots, ...]` allocations, passed unsliced. Certified at 3..64 - and 461 slots; at 461 the result is bitwise identical to the 64-slot - result for the same slot content. + `[layers, slots, ...]` allocations, passed unsliced. Certified at 4, 6 + and 461 slots, over an fp32 and a bf16 recurrent pool. Loud, from the callable itself (ValueError): an empty recurrent pool, conv and recurrent slot counts differing (the in-range test uses the recurrent diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/patch_embed_conv3d.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/patch_embed_conv3d.py index 9a6fa3653b37..561418cadd79 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/patch_embed_conv3d.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/patch_embed_conv3d.py @@ -47,63 +47,19 @@ class PatchEmbedConv3d(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec: num_patches, dtype; `in_channels`, `kernel`, `embed_dim` (default - # the checkpoint's); `bias` (default True); `layout` "flat" (default), - # "five_d" ([N, C, kt, kh, kw]), "column_slice" (a [N, 2*row] buffer's first - # half) or "transposed" (column-major rows). fp32 cells run with - # `torch.backends.cudnn.allow_tf32 = False` -- see note. - CELLS: tuple[Cell, ...] = ( - tuple( - Cell( - why=f"checkpoint geometry, {n} patches bf16", - spec=dict(num_patches=n, dtype=torch.bfloat16), - ) - for n in (1, 2, 7, 64, 1024, 4096) - ) - + ( - Cell( - why="tower call ceiling: 8192 tokens x merge 2x2 = 32768 patches bf16", - spec=dict(num_patches=32768, dtype=torch.bfloat16), - ), - ) - + ( - Cell( - why="explicit [N, C, kt, kh, kw] rows", - spec=dict(num_patches=64, dtype=torch.bfloat16, layout="five_d"), - ), - Cell( - why="bias=None (the Qwen2-VL form)", - spec=dict(num_patches=128, dtype=torch.bfloat16, bias=False), - ), - Cell(why="fp16", spec=dict(num_patches=256, dtype=torch.float16)), - Cell( - why="fp32 with cudnn TF32 off", - spec=dict(num_patches=256, dtype=torch.float32), - ), - ) - + tuple( - Cell( - why=f"geometry read from weight: C={c} kernel={k} out={o}", - spec=dict( - num_patches=33, - dtype=torch.bfloat16, - in_channels=c, - kernel=k, - embed_dim=o, - ), - ) - for c, k, o in ((1, (1, 4, 4), 8), (4, (2, 8, 8), 64), (3, (1, 16, 16), 128)) - ) - + ( - Cell( - why="column slice of a wider pixel buffer", - spec=dict(num_patches=8, dtype=torch.bfloat16, layout="column_slice"), - ), - Cell( - why="column-major pixel buffer (row elements strided)", - spec=dict(num_patches=8, dtype=torch.bfloat16, layout="transposed"), - ), + # spec: num_patches, dtype. The checkpoint's geometry, flat contiguous + # rows and a bias: what the vision tower passes. + CELLS: tuple[Cell, ...] = tuple( + Cell( + why=f"checkpoint geometry, {n} patches bf16", + spec=dict(num_patches=n, dtype=torch.bfloat16), ) + for n in (1, 64, 4096) + ) + ( + Cell( + why="tower call ceiling: 8192 tokens x merge 2x2 = 32768 patches bf16", + spec=dict(num_patches=32768, dtype=torch.bfloat16), + ), ) note = """ @@ -125,8 +81,10 @@ class PatchEmbedConv3d(OpWrapper): depend on how many rows share the call. Check rows against `reference`, never against a call of a different size. - Not certified: float8 operands; patch counts between the 4096- and 32768-row - cells, and past 32768. + Not certified: dtypes other than bf16 (the fp32 TF32 behavior above is + measured, not certified); other geometries, `bias=None` and non-contiguous + pixel rows; patch counts between the 4096- and 32768-row cells, and past + 32768. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/pos_embed_interpolate.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/pos_embed_interpolate.py index ba26e3250f11..5fc3075f782b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/pos_embed_interpolate.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/pos_embed_interpolate.py @@ -28,19 +28,12 @@ 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. +#: Grids the processor produces: the smallest, the exact-fit grid, the +#: non-square serving grid and the largest one call carries (32768 patches). 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),), ) @@ -74,59 +67,35 @@ class PosEmbedInterpolate(OpWrapper): 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 + 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 + ), ) - + ( - 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, - ), + 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, - ), + ), + 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 = """ @@ -147,7 +116,8 @@ class PosEmbedInterpolate(OpWrapper): at capture. Cost is ~225 us of launch-bound index construction per image, flat in token count; the gather is the only part that scales. - Not certified: a CPU table (it runs and returns on CPU). + Not certified: a CPU table (it runs and returns on CPU); fp16 and fp32 + tables; a `spatial_merge_size` other than the checkpoint's 2. """ def __call__( diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 6251b98757b7..87c73a700e44 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -71,7 +71,7 @@ l0_cpu: # tensors and the CPU block_scale_interleave. The qwen4_exp entries are the target's # tests that run on the source text, stubbed weights and meta tensors: its multimodal # splice and admission check on a stub tower, the KV pool contract, its KV storage - # derivation and attention-call arguments, and its routed-expert scalars. + # derivation and attention-call arguments, its routed-expert scalars and its weight table. # The rest of unittest/_torch/modeling_v2 needs a GPU. - unittest/_torch/modeling_v2/test_modeling_v2_claims.py - unittest/_torch/modeling_v2/test_modeling_v2_routing.py @@ -84,6 +84,7 @@ l0_cpu: - unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_pool_contract.py - unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_kv_quant.py - unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_thop_attention_calls.py + - unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_weight_table.py - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_fp8_indexer_rejects_different_qk_norm_epsilons - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_moe_reduces_only_local_terms - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_decoder_layer_sets_post_fusion_from_moe_scheduler diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py index 0bd6a6662bbe..71bc78290455 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py @@ -22,81 +22,57 @@ def _seed(text: str) -> int: def _build(spec, seed): - """`(buffer, out, region, gate_buffer, gate)` for one spec. + """`(out, gate_buffer, gate)` for one spec. - `out` is `region(buffer)`: the whole buffer, a padded window of it, or every - other row. `gate` is a view of `gate_buffer` -- for the checkpoint layout the - second half of each `[T, heads, 2 * D]` q_proj head. + `gate` is a view of `gate_buffer`: the second half of each + `[T, heads, 2 * D]` q_proj head, or the 2-D gate window of the target's + fused projection row. """ g = torch.Generator(device="cuda").manual_seed(seed) t, heads, d, dtype = spec["tokens"], spec["heads"], spec["head_dim"], spec["dtype"] hidden = heads * d - def randn(*shape, dt=dtype): - return torch.randn(*shape, generator=g, device="cuda").to(dt) - - layout = spec.get("out", "contiguous") - if layout == "window": - buffer = randn(t + 2, hidden + 2048) - - def region(x): - return x[:t, :hidden] - elif layout == "every_other_row": - buffer = randn(2 * t, hidden) - - def region(x): - return x[::2] - else: - buffer = randn(t, hidden) - - def region(x): - return x + def randn(*shape): + return torch.randn(*shape, generator=g, device="cuda").to(dtype) - gate_layout = spec["gate"] - if gate_layout == "interleaved": + out = randn(t, hidden) + if spec["gate"] == "interleaved": gate_buffer = randn(t, heads, 2 * d) gate = gate_buffer[:, :, d:] - elif gate_layout == "packed": - gate_buffer = gate = randn(t, heads, d) - elif gate_layout == "fused_row": + else: + assert spec["gate"] == "fused_row", spec["gate"] # The target's fused input projection row [q | k | v | gate | index | pad]: # 14080 columns, the gate a 2-D [T, hidden] window at column 7168. gate_buffer = randn(t, 14080) gate = gate_buffer[:, 7168 : 7168 + hidden] - else: - gate_buffer = gate = randn(t, hidden, dt=spec.get("gate_dtype", dtype)) - return buffer, region(buffer), region, gate_buffer, gate + return out, gate_buffer, gate @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name certifying(op) - buffer, out, region, gate_buffer, gate = _build(cell.spec, seed=_seed(cell.why)) + out, 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() + gate_before = gate_buffer.clone() expected = op.reference(out, gate) result = op(out, gate) assert result.data_ptr() == out.data_ptr() and result.shape == out.shape torch.testing.assert_close(result, expected) - # In place: everything outside the [T, hidden] window, and the gate, bitwise unchanged. - touched = torch.zeros_like(buffer, dtype=torch.bool) - region(touched).fill_(True) - assert torch.equal(buffer[~touched], buffer_before[~touched]) assert torch.equal(gate_buffer, gate_before) def test_the_gate_half_is_read_through_its_own_strides() -> None: """The strided view and its compacted copy agree bitwise; the q half of the same buffer does not -- so the cells can tell the two halves apart.""" - spec = dict(op.CELLS[1].spec) - _, out, _, q_and_gate, gate = _build(spec, seed=7) + spec = dict(op.CELLS[0].spec, tokens=17) + out, q_and_gate, gate = _build(spec, seed=7) q_half = q_and_gate[:, :, : spec["head_dim"]] from_strided = op(out.clone(), gate) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_erf.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_erf.py index ab0b602f1189..9beeec3d0c99 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_erf.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_erf.py @@ -67,18 +67,7 @@ def _build(spec, seed): ) bias = FC1_BIAS_STD * torch.randn(MERGER_WIDTH, generator=g, device="cuda") return (grouped @ weight.to(dtype).float().t() + bias.to(dtype).float()).to(dtype) - layout = spec.get("layout", "contiguous") lo, hi = spec["lo"], spec["hi"] - if layout == "column_slice": - rows, width = shape - x = _uniform((rows, 2 * width), lo, hi, dtype, seed)[:, :width] - assert not x.is_contiguous() and x.stride(-1) == 1 - return x - if layout == "transposed": - rows, width = shape - x = _uniform((width, rows), lo, hi, dtype, seed).t() - assert x.stride(-1) != 1 - return x return _uniform(shape, lo, hi, dtype, seed) @@ -89,15 +78,8 @@ def test_certified_cells(cell, certifying): # noqa: F811 before = x.clone() out = op(x) assert out.shape == x.shape and out.dtype == x.dtype and out.device == x.device - if x.numel(): - assert out.data_ptr() != x.data_ptr(), "the entry is out-of-place" + assert out.data_ptr() != x.data_ptr(), "the entry is out-of-place" assert torch.equal(before, x), "the entry must not mutate its input" - if cell.spec.get("layout") == "transposed": - assert out.stride() == x.stride(), "the output keeps the input's memory format" - assert torch.equal(out, op(x.contiguous())) - if cell.spec.get("layout") == "column_slice": - assert out.is_contiguous() - assert torch.equal(out, op(x.contiguous())) torch.testing.assert_close(out, op.reference(x)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_tanh.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_tanh.py index aa52e550fce0..c8a7fa602db3 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_tanh.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_tanh.py @@ -58,18 +58,7 @@ def _build(spec, seed): ).to(dtype) bias = (FC1_BIAS_STD * torch.randn(shape[-1], generator=g, device="cuda")).to(dtype) return (hidden.float() @ weight.float().t() + bias.float()).to(dtype) - layout = spec.get("layout", "contiguous") lo, hi = spec["lo"], spec["hi"] - if layout == "column_slice": - rows, width = shape - x = _uniform((rows, 2 * width), lo, hi, dtype, seed)[:, :width] - assert not x.is_contiguous() and x.stride(-1) == 1 - return x - if layout == "transposed": - rows, width = shape - x = _uniform((width, rows), lo, hi, dtype, seed).t() - assert x.stride(-1) != 1 - return x return _uniform(shape, lo, hi, dtype, seed) @@ -80,15 +69,8 @@ def test_certified_cells(cell, certifying): # noqa: F811 before = x.clone() out = op(x) assert out.shape == x.shape and out.dtype == x.dtype and out.device == x.device - if x.numel(): - assert out.data_ptr() != x.data_ptr(), "the entry is out-of-place" + assert out.data_ptr() != x.data_ptr(), "the entry is out-of-place" assert torch.equal(before, x), "the entry must not mutate its input" - if cell.spec.get("layout") == "transposed": - assert out.stride() == x.stride(), "the output keeps the input's memory format" - assert torch.equal(out, op(x.contiguous())) - if cell.spec.get("layout") == "column_slice": - assert out.is_contiguous() - assert torch.equal(out, op(x.contiguous())) torch.testing.assert_close(out, op.reference(x)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_silu_and_mul_rows.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_silu_and_mul_rows.py index 7502b3963747..2fe01861e216 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_silu_and_mul_rows.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_silu_and_mul_rows.py @@ -26,7 +26,7 @@ def _build(spec, seed): g = torch.Generator(device="cuda").manual_seed(seed) - rows, width = spec["rows"], spec.get("width", 2 * INTER + 8) + rows, width = spec["rows"], 2 * INTER + 8 x = (torch.randn(rows, width, generator=g, device="cuda") * 3.0).to(torch.bfloat16) if rows: x[0, :4] = torch.tensor([0.0, -0.0, 88.0, -100.0], device="cuda").to(torch.bfloat16) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_canonical_selection.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_canonical_selection.py index 3d25d1ad2efd..46b72e2c21fe 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_canonical_selection.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_canonical_selection.py @@ -55,10 +55,6 @@ def _logits(kind: str, rows: int, cols: int, gen: torch.Generator) -> torch.Tens x = _logits("relu_sparse", rows, cols, gen) negative = torch.rand(rows, cols, generator=gen, device=DEV) < 0.5 return torch.where((x == 0) & negative, torch.full_like(x, -0.0), x) - if kind == "two_valued": - return torch.randint(0, 2, (rows, cols), generator=gen, device=DEV).float() - if kind == "all_equal": - return torch.full((rows, cols), 1.5, device=DEV) if kind == "pair_at_k": x = torch.rand(rows, cols, generator=gen, device=DEV).argsort(1).float() + 1.0 order = torch.argsort(x[:, :2048], dim=1, descending=True) @@ -78,11 +74,9 @@ def _windows(spec: dict, rows: int, cols: int, gen: torch.Generator): lo, hi = spec["windows"][1:] positions = torch.randint(lo, hi + 1, (rows,), generator=gen, device=DEV) ends = ((positions + 1) // COMPRESS_RATIO).clamp_max(cols).to(torch.int32) - elif kind == "ends": - ends = torch.tensor(spec["windows"][1], dtype=torch.int32, device=DEV) else: - starts = torch.tensor(spec["windows"][1], dtype=torch.int32, device=DEV) - ends = torch.tensor(spec["windows"][2], dtype=torch.int32, device=DEV) + assert kind == "ends", kind + ends = torch.tensor(spec["windows"][1], dtype=torch.int32, device=DEV) assert starts.numel() == rows and ends.numel() == rows return starts, ends @@ -218,19 +212,16 @@ def test_narrowing_the_table_to_a_bound_on_the_windows_changes_nothing() -> None assert torch.equal(wide, narrow) -@pytest.mark.parametrize("kind", ["relu_sparse", "two_valued", "eighths", "pair_at_k"]) +@pytest.mark.parametrize("kind", ["relu_sparse", "pair_at_k"]) def test_radix_launches_differ_and_canonical_rows_do_not(kind) -> None: """Repeated radix launches over one tie-heavy score matrix. The raw rows differ from launch to launch -- in order, and where the tie class at the k-th score straddles the budget, in which tied columns are kept -- and the canonical row is one and the same - row, the reference's, for every launch of every row.""" + row, the reference's, for every launch of every row. Zero ties as the target's relu + scores have them, and a positive pair tied exactly at the k-th value.""" gen = torch.Generator(device=DEV).manual_seed(17) rows, launches = 64, 20 - if kind == "eighths": - # Positive ties: 64 levels over 2048 columns. - logits = torch.randint(0, 64, (rows, COLS), generator=gen, device=DEV).float() / 8 - else: - logits = _logits(kind, rows, COLS, gen) + logits = _logits(kind, rows, COLS, gen) starts = torch.zeros(rows, dtype=torch.int32, device=DEV) ends = torch.full((rows,), 2048, dtype=torch.int32, device=DEV) _poison(logits, starts, ends, "nan") diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_expand_block_indices.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_expand_block_indices.py index 046684168aec..0904d0565654 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_expand_block_indices.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_expand_block_indices.py @@ -100,18 +100,12 @@ def _build(spec, seed: int): 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 = torch.randint(0, 64, (rows + 37, block_topk), dtype=torch.int32, device=DEV) 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) @@ -168,10 +162,11 @@ def test_launch_shapes_agree_on_shared_rows() -> None: @pytest.mark.parametrize( "position, expected_tail", - [(4095, 0), (4096, 1), (4097, 2), (4098, 3), (4099, 0), (7, 0), (3, 0), (0, 1)], + [(4098, 3), (4099, 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.""" + """`(position + 1) % cr` tokens are appended right after the selection, ascending: the + widest tail, none, and the first row. Every width 0..3 is a certified cell.""" 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 @@ -260,22 +255,19 @@ def test_output_is_fresh_and_inputs_are_untouched() -> None: def test_op_rejects_malformed_arguments() -> None: - """Every metadata violation is the callable's own loud error: why is_valid is empty.""" + """Every metadata violation is the callable's own loud error: why is_valid is empty. + One case per error family.""" 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", ( @@ -291,7 +283,6 @@ def test_op_rejects_malformed_arguments() -> None: {}, ), ("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): diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_mtp_reuse_topk.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_mtp_reuse_topk.py index 538d345d288d..a8746a2040e2 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_mtp_reuse_topk.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_mtp_reuse_topk.py @@ -362,11 +362,6 @@ def strided(tensor): if name == "output column slice": kwargs["output"] = torch.zeros((rows, 2 * K), dtype=torch.int32, device=DEV)[:, :K] - elif name == "shared_indices column slice": - wide = torch.full((kwargs["shared_indices"].shape[0], 2 * K), -1, dtype=torch.int32) - wide = wide.to(DEV) - wide[:, :K] = kwargs["shared_indices"] - kwargs["shared_indices"] = wide[:, :K] elif name == "shared_indices 3-D, middle extent == block_topk": # Slot 1 meant; the 3-D buffer passes the width check and reads slot 0. stack = torch.full((2, 8, 8), -1, dtype=torch.int32, device=DEV) @@ -391,14 +386,13 @@ def strided(tensor): return intended, kwargs +#: One case per guard family: a non-contiguous output, strided row metadata, and a +#: whole multi-slot buffer passed where one slot's view belongs (indices and counts). GUARDS = ( ("output column slice", "output must be contiguous"), - ("shared_indices column slice", "shared_indices must be one slot's"), ("shared_indices 3-D, middle extent == block_topk", "shared_indices must be one slot's"), ("shared_visible_blocks 2-D [slots, seqs]", "shared_visible_blocks must be one slot's"), ("strided request_indices", "request_indices must be contiguous"), - ("strided visible_blocks", "visible_blocks must be contiguous"), - ("strided shared_visible_blocks", "shared_visible_blocks must be one slot's"), ) @@ -415,17 +409,16 @@ def test_guard_refuses_a_silently_wrong_layout(name: str, message: str) -> None: def test_op_rejects_malformed_arguments() -> None: - """The callable's own loud errors: why `is_valid` does not restate them.""" - kwargs = _build(_spec_of("dense standalone"), seed=17)["kwargs"] + """The callable's own loud errors, one per family: why `is_valid` does not restate them.""" + spec = {**_spec_of("8 requests, one per boundary"), "layout": "dense"} + kwargs = _build(spec, seed=17)["kwargs"] cases = [ ("width does not match", dict(shared_indices=kwargs["shared_indices"][:, : K // 2])), ( "one request and bound per query row", dict(request_indices=kwargs["request_indices"][:3]), ), - ("one request and bound per query row", dict(visible_blocks=kwargs["visible_blocks"][:3])), ("requires CUDA tensors", dict(shared_indices=kwargs["shared_indices"].cpu())), - ("requires CUDA tensors", dict(output=kwargs["output"].cpu())), ] for message, override in cases: with pytest.raises(ValueError, match=message): diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_scores.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_scores.py index 89e2c0ad7711..62aa3a5b936f 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_scores.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_scores.py @@ -38,7 +38,7 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: 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 + (asserted at 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) @@ -130,23 +130,14 @@ def _rows_and_requests(spec, gen): def pages_for(seq_len): return math.ceil(seq_len / tpb) - if layout in ("random", "shuffled", "past_table", "arange"): + if layout in ("random", "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": + req = torch.arange(rows, device=DEV) * requests // max(rows, 1) + if 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": @@ -205,10 +196,6 @@ def _build(spec, seed): # [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 @@ -293,39 +280,31 @@ def _capture_args(seed, only_visible_blocks=False): return _build(spec, seed) -@pytest.mark.parametrize("shape", ["512rows_narrow", "8191rows_capture"]) -def test_the_two_kernels_agree(shape) -> None: +def test_the_two_kernels_agree() -> 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. + and again at 8191, so the packed kernel does not drift with batch size; + asserted at the 8191-row capture shape. """ - if shape == "512rows_narrow": - args = _build(_case(512, 6), seed=404) - else: - args = _capture_args(seed=8700) + 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: +@pytest.mark.parametrize("context_rows", [False, True], ids=["row_wise", "query_tile"]) +def test_results_are_bitwise_reproducible(context_rows) -> 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) + whole tensor is written; once per kernel.""" + args = _build(_case(300, 5, context_rows=context_rows), 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"] -) +@pytest.mark.parametrize("shape", ["64rows_row_wise", "8191rows_capture"]) def test_only_visible_blocks_leaves_the_tail_unspecified(shape) -> None: """`only_visible_blocks=True` writes the visible columns and nothing else. @@ -337,9 +316,7 @@ def test_only_visible_blocks_leaves_the_tail_unspecified(shape) -> None: 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) + args = _build(_case(64, 4, only_visible_blocks=True), seed=670) expected = op.reference(**args) poison = torch.full(expected.shape, 98765.0, dtype=torch.float32, device=DEV) del poison @@ -492,21 +469,18 @@ def scores(*, offset=ratio - 1, relu=True, scale=head_dim**-0.5, mean=False, no_ } -@pytest.mark.parametrize("width", ["8pages", "260pages"]) -def test_wrong_variants_fail_the_gate(width) -> None: +def test_wrong_variants_fail_the_gate() -> 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. + At the target's 260-page width, 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 - ) + 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) @@ -580,14 +554,6 @@ def _mutated(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", @@ -596,40 +562,23 @@ def _mutated(args, **overrides): 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`.""" + entry has no `is_valid`. One case per refusal family: operand shape, block + table, dtypes, geometry scalars, metadata.""" mutate, message = _REJECTIONS[case] args = mutate(_build(_case(16, 2), seed=1701)) with pytest.raises(ValueError, match=message): diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_store.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_store.py index 151d027caab5..85b8e69044c4 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_store.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_store.py @@ -51,25 +51,14 @@ def _bits(tensor: torch.Tensor) -> torch.Tensor: return tensor.contiguous().view(as_int) -def _hole(spec, table, live_requests, allocated_cols): - """One row that resolves to no slot, and the table edit it needs.""" - tpb, width, slots, pages = (spec[k] for k in ("tokens_per_block", "width", "slots", "pages")) +def _hole(spec, table, allocated_cols): + """One row that resolves to no slot: its (request, position).""" 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 + return 0, allocated_cols[0] * spec["tokens_per_block"] + assert kind == "negative_request" + return -1, 0 def _shared_page_planes(spec, extras: dict): @@ -115,19 +104,11 @@ def _build(spec, seed: int): allocated_cols.append(cols) assert used <= pages and len(segments) < slots, "pool or table too small for the batch" if spec["hole"] is not None: - request, position = _hole(spec, table, len(segments), allocated_cols) + request, position = _hole(spec, table, allocated_cols) requests.append(request) positions.append(position) rows += 1 table = table.to(DEV) - if spec["table_layout"] == "column_slice": - wide = torch.full((slots, width + 4), -1, dtype=torch.int32, device=DEV) - wide[:, 2 : 2 + width] = table - table = wide[:, 2 : 2 + width] - assert not table.is_contiguous() - elif spec["table_layout"] == "transposed": - table = table.t().contiguous().t() - assert table.stride(1) != 1 projection = torch.randn( rows, spec["index_heads"] + hkv, head_dim, generator=gen, dtype=torch.float32 @@ -138,43 +119,23 @@ def _build(spec, seed: int): # pad] = 110 heads of 128: the raw index key is head 108. row = torch.randn(rows, 110, head_dim, generator=gen, dtype=torch.float32) token_k = row.to(device=DEV, dtype=spec["key_dtype"])[:, 108 : 108 + hkv] - 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() + # The engine's `[3, max_num_tokens]` position buffer, transposed. + buffer = torch.full((3, ENGINE_MAX_NUM_TOKENS), -1, dtype=torch.int32, device=DEV) + buffer[:, :rows] = coords.t() + coords = buffer[:, :rows].t() extras = {} position_cache = None if spec["cache_layout"] == "shared_page": index_cache, position_cache = _shared_page_planes(spec, extras) - elif 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 @@ -198,29 +159,25 @@ def _build(spec, seed: int): def test_certified_cells(cell, certifying): # noqa: F811 certifying(op) args, extras = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) - before = args["position_cache"].clone() ref = op.reference(**args) out = op(**args) assert out[0] is args["index_cache"] and out[1] is args["position_cache"] _compare(out, ref) - if "neighbour" in extras: - view, saved = extras["neighbour"] - assert torch.equal(_bits(view), _bits(saved)), "the neighbouring role moved" if "roles" in extras: fill = int(_bits(torch.tensor([KEY_FILL], dtype=torch.bfloat16))[0]) assert bool((extras["roles"].view(torch.int16) == fill).all()), ( "another index-K role of the shared page moved" ) - 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]) + + +#: The cell of 32 MTP7 verify blocks (256 rows): the segments, pool and table it runs over. +_MTP7_32 = next(c.spec for c in op.CELLS if c.why.startswith("decode MTP7: 32 requests")) +_MTP7_BATCH = {key: _MTP7_32[key] for key in ("segments", "pages", "slots")} def _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} + spec = {**_MTP7_32, "segments": ((0, 1), (7, 1), (32, 1), (40, 1)), "pages": 8} return _build({**spec, **overrides}, seed=seed)[0] @@ -279,15 +236,10 @@ def test_repeats_are_stream_ordered_and_idempotent() -> None: assert torch.equal(args["position_cache"], first[1]) -#: The cell of 32 MTP7 verify blocks (256 rows): the segments, pool and table it runs over. -_MTP7_32 = next(c.spec for c in op.CELLS if c.why.startswith("decode MTP7: 32 requests")) -_MTP7_BATCH = {key: _MTP7_32[key] for key in ("segments", "pages", "slots")} - - -@pytest.mark.parametrize("batch", ({}, _MTP7_BATCH), ids=("decode", "mtp7-32x8")) -def test_cuda_graph_capture_and_replay(batch) -> None: - """Captured once, replayed against new positions written into the same buffer.""" - args = _small(seed=52, **batch) +def test_cuda_graph_capture_and_replay() -> None: + """Captured once at the 32 MTP7 verify blocks (256 rows), replayed against new + positions written into the same buffer.""" + args = _small(seed=52, **_MTP7_BATCH) positions = args["logical_positions"] replayed = positions.clone() positions.zero_() @@ -336,6 +288,8 @@ def test_the_stored_key_is_raw_not_normalized() -> None: # # Each is a shape the kernel would otherwise read past or write wrongly; the # entry refuses it before launching, so `is_valid` has nothing left to guard. +# One case per refusal family: metadata, row counts, the index plane, the +# position plane, the table. def _rejected(args, match: str) -> None: @@ -346,60 +300,30 @@ def _rejected(args, match: str) -> None: assert torch.equal(args["position_cache"], before[1]) -def test_rejects_two_dimensional_request_indices() -> None: - args = _small(seed=60) - args["request_indices"] = args["request_indices"][:, None] - _rejected(args, "one-dimensional") - - def test_rejects_metadata_of_unequal_length() -> None: args = _small(seed=61) args["logical_positions"] = args["logical_positions"][:1] _rejected(args, "equally long") -def test_rejects_float_metadata() -> None: - args = _small(seed=62) - args["logical_positions"] = args["logical_positions"].float() - _rejected(args, "int32 or int64") - - def test_rejects_a_single_token_k_row() -> None: args = _small(seed=63) args["token_k"] = args["token_k"][:1] _rejected(args, "token_k") -def test_rejects_a_single_coordinate_row() -> None: - args = _small(seed=64) - args["position_coordinates"] = args["position_coordinates"][:1] - _rejected(args, "position_coordinates") - - def test_rejects_a_page_size_that_is_not_the_cache_s() -> None: args = _small(seed=65) args["tokens_per_block"] = 16 _rejected(args, "index_cache") -def test_rejects_a_position_plane_with_a_different_page_size() -> None: - args = _small(seed=66) - args["position_cache"] = torch.full((8, 16, 3), POS_FILL, dtype=torch.int32, device=DEV) - _rejected(args, "position_cache") - - def test_rejects_a_position_plane_narrower_than_the_coordinates() -> None: args = _small(seed=67) args["position_cache"] = torch.full((8, 32, 2), POS_FILL, dtype=torch.int32, device=DEV) _rejected(args, "position_cache") -def test_rejects_a_cache_with_more_heads_than_the_key() -> None: - args = _small(seed=68) - args["index_cache"] = torch.full((8, 32, 2, 128), KEY_FILL, device=DEV).to(torch.bfloat16) - _rejected(args, "index_cache") - - def test_rejects_a_one_dimensional_table() -> None: args = _small(seed=69) args["block_table"] = args["block_table"][0] diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store.py index 5fed6e17c22a..3a906789e826 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store.py @@ -178,10 +178,10 @@ def test_fp8_pool_saturates_an_out_of_range_row() -> None: _MTP7_BATCH = {key: _MTP7_32[key] for key in ("segments", "pages", "slots")} -@pytest.mark.parametrize("batch", ({}, _MTP7_BATCH), ids=("decode", "mtp7-32x8")) -def test_cuda_graph_replay_follows_updated_positions(batch) -> None: - """Captured once, replayed against new positions written into the same buffer.""" - args, pool = _decode_args(seed=9, **batch) +def test_cuda_graph_replay_follows_updated_positions() -> None: + """Captured once at 32 MTP7 verify blocks, replayed against new positions written + into the same buffer.""" + args, pool = _decode_args(seed=9, **_MTP7_BATCH) positions = args["logical_positions"] saved = positions.clone() positions.zero_() diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store_nvfp4.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store_nvfp4.py index 6c934acfed42..8d939576ef43 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store_nvfp4.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store_nvfp4.py @@ -427,7 +427,7 @@ def test_rounding_rivals_are_rejected() -> None: assert moved > 0, "ties away from zero must move some K bytes on every-finite-bf16 rows" # V's scale layout, on random rows (the every-finite rows saturate V's scales). - spec = next(c.spec for c in op.CELLS if c.why.startswith("mixed:")) + spec = next(c.spec for c in op.CELLS if c.why.startswith("prefill:")) args, _, _ = _build(spec, seed=22) out = op(**args) page = int(args["block_table"][0, 0]) @@ -441,31 +441,6 @@ def test_rounding_rivals_are_rejected() -> None: assert not torch.equal(interleaved, token_major), "the V scale layouts must be told apart" -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): - call() - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - call() - pool[0].fill_(DATA_FILL) - pool[1].fill_(SCALE_FILL) - positions.copy_(saved) - ref = op.reference(**args) - graph.replay() - torch.cuda.synchronize() - _compare((args["k_cache"], args["v_cache"], args["k_scale_cache"], args["v_scale_cache"]), ref) - - def test_op_rejects_operands_it_cannot_serve() -> None: """The callable's own loud errors: why the entry carries no `is_valid`.""" args, _, _ = _decode_args(seed=5) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa.py index ad2dd5c7d32a..b44ab112e3dc 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa.py @@ -384,7 +384,7 @@ def test_calls_are_bitwise_deterministic() -> None: assert torch.equal(op(**args), op(**args)) -@pytest.mark.parametrize("rows", [1, 4, 128]) +@pytest.mark.parametrize("rows", [1, 128]) def test_output_gate_in_the_merge_is_the_standalone_gate(rows) -> None: """The gate carried by the merge equals the ungated call followed by the standalone gate kernel, bit for bit, and the gate moves the output.""" @@ -467,34 +467,14 @@ def test_gate_discriminates_wrong_variants() -> None: _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", @@ -518,7 +498,9 @@ def test_gate_discriminates_wrong_variants() -> None: @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.""" + One per guard family: operand dtype, metadata dtype, scalar value, head + geometry, the prefill pair and the device. 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] diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_decode.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_decode.py index bf9206712e68..99060ead085b 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_decode.py @@ -96,13 +96,13 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0) -@pytest.mark.parametrize("gated", [True, False]) -def test_matches_the_composition(gated) -> None: - args = _build(dict(pos=(8700,), gated=gated), seed=91) +def test_matches_the_composition() -> None: + """One representative gated decode row; the composition is certified on its own.""" + args = _build(dict(pos=(8700,), gated=True), seed=91) got = op(*args) want = qsa_paged_sparse_gqa(**_kw(args)) assert_within_ulp(got, want, element_ulp=4.0, rms_ulp=2.0) - print(f"\n gated={gated}: bitwise the composition's {(got == want).float().mean().item():.4f}") + print(f"\n bitwise the composition's {(got == want).float().mean().item():.4f}") def _replays_like_eager(step, refresh) -> None: diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4.py index 8c3c4a31e8a8..fcb9ab66568e 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4.py @@ -385,18 +385,10 @@ def test_zero_rows_return_an_empty_output() -> None: lambda a: {**a, "k_cache": a["k_cache"].view(torch.int8)}, "k_cache must be uint8", ), - "scales_of_another_width": ( - lambda a: {**a, "v_scale_cache": a["v_scale_cache"][..., :8]}, - "v_scale_cache must be uint8", - ), "single_scale": ( lambda a: {**a, "kv_scale_quant_orig": a["kv_scale_quant_orig"][:1]}, "kv_scale_quant_orig must be a contiguous fp32 [3]", ), - "int64_block_table": ( - lambda a: {**a, "block_table": a["block_table"].long()}, - "block_table must be a row-contiguous 2-D int32 tensor", - ), "positions_without_ratio": ( lambda a: {**a, "query_positions": torch.zeros(4, dtype=torch.int64, device=DEV)}, "pass both query_positions and compress_ratio, or neither", @@ -407,7 +399,9 @@ def test_zero_rows_return_an_empty_output() -> None: @pytest.mark.parametrize("case", list(_REJECTIONS), ids=list(_REJECTIONS)) def test_rejects(case) -> None: - """The op's own loud refusals -- the reason the entry carries no `is_valid`.""" + """The op's own loud refusals -- the reason the entry carries no `is_valid`. + One per guard family: q's dtype, the pool planes, the second-level scales, + the prefill pair and the device.""" args, _ = _build(_with(segments=((600, 1),) * 4), seed=9100) args["selected_tokens"] = args["selected_tokens"][:, :64].contiguous() mutate, message = _REJECTIONS[case] diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_prefill_compress.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_prefill_compress.py index 0999b98a150c..c1b5a3dee813 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_prefill_compress.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_prefill_compress.py @@ -211,9 +211,10 @@ def test_a_second_call_over_the_same_rows_is_not_idempotent() -> None: _MTP7_BATCH = {key: _MTP7_32[key] for key in ("segments", "pages", "slots", "rounding_flips")} -@pytest.mark.parametrize("batch", ({}, _MTP7_BATCH), ids=("prefill", "mtp7-32x8")) -def test_cuda_graph_capture_and_replay(batch) -> None: - """The launch grid is fixed by the row count, so it captures and replays.""" +def test_cuda_graph_capture_and_replay() -> None: + """The launch grid is fixed by the row count, so it captures and replays: 32 MTP7 + verify blocks, the decode target's captured batch.""" + batch = _MTP7_BATCH args = _scene(seed=16, **batch) cache = args["index_cache"] before = cache.clone() @@ -234,7 +235,7 @@ def test_cuda_graph_capture_and_replay(batch) -> None: cache.copy_(before) graph.replay() torch.cuda.synchronize() - _compare(cache, ref, batch.get("rounding_flips", 0)) + _compare(cache, ref, batch["rounding_flips"]) def test_stream_ordering_needs_no_explicit_synchronization() -> None: @@ -310,16 +311,16 @@ def _strided(tensor: torch.Tensor, dim: int) -> torch.Tensor: "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. + One per guard form: a 1-D operand read flat (`request_indices` and + `k_norm_weight` carry the same `is_contiguous` guard), and `cos_sin`'s last + dim. Old receipt's max abs error against the right answer: 4.72 and 7.16 + respectively; the floor below states a bound, not one draw. """ args = _scene(seed=12) before = args["index_cache"].clone() @@ -358,9 +359,7 @@ def _variant(args, defect): 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) + normed = (pooled * scale * (args["k_norm_weight"].float() + 1.0)).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 @@ -375,13 +374,8 @@ def _variant(args, defect): 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) + first, second = normed[:, :n], normed[:, n : 2 * 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 ) @@ -392,14 +386,12 @@ def _variant(args, defect): "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. + """Every defect differs in at least one bit; two 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. diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py index 5bfb34d165cc..e7eaac4ba272 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py @@ -143,14 +143,9 @@ def _build(spec, seed): pool = torch.randn(pool_pages, 2, kv, tpb, d, generator=gen, device=DEV).to(spec["cache_dtype"]) k_cache, v_cache = pool[:, 0], pool[:, 1] - if spec["cache_layout"] == "contiguous": - k_cache, v_cache = k_cache.contiguous(), v_cache.contiguous() - q_dtype = spec["q_dtype"] - if spec["q_layout"] == "qkv_slice": - qkv = torch.randn(rows, q_heads + 2 * kv, d, generator=gen, device=DEV).to(q_dtype) - q = qkv[:, :q_heads] - else: - q = torch.randn(rows, q_heads, d, generator=gen, device=DEV).to(q_dtype) + assert spec["cache_layout"] == "pool_planes" and spec["q_layout"] == "qkv_slice" + qkv = torch.randn(rows, q_heads + 2 * kv, d, generator=gen, device=DEV).to(spec["q_dtype"]) + q = qkv[:, :q_heads] selected = _selected_tokens(positions, seed + 3, spec["compress_ratio"], spec["token_topk"], spec.get("shared_scores", False)) # fmt: skip @@ -310,16 +305,13 @@ def test_zero_rows_return_an_empty_output() -> None: assert out.shape == (0, CKPT["q_heads"], CKPT["head_dim"]) and out.dtype == torch.bfloat16 +#: One loud rejection per family: the missing prefill pair, a dtype, a shape and a layout. _REJECTIONS = { "no_prefill_pair": ( lambda a: {**a, "query_positions": None, "compress_ratio": None}, "the row-tile kernel needs the prefill pair", ), "float32_query": (lambda a: {**a, "q": a["q"].float()}, "QSA queries must use BF16"), - "int64_selection": ( - lambda a: {**a, "selected_tokens": a["selected_tokens"].long()}, - "QSA selected tokens must be a two-dimensional int32 tensor", - ), "width_vs_ratio": ( lambda a: {**a, "selected_tokens": a["selected_tokens"][:, :64].contiguous()}, "incompatible with its compression ratio", @@ -328,7 +320,6 @@ def test_zero_rows_return_an_empty_output() -> None: lambda a: {**a, "request_indices": torch.zeros(8, dtype=torch.int32, device=DEV)[::2]}, "QSA request indices must be contiguous", ), - "host_query": (lambda a: {**a, "q": a["q"].cpu()}, "needs every tensor on one CUDA device"), } diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_expand_decode.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_expand_decode.py index 966013ce94d0..84ff70dd08e9 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_expand_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_expand_decode.py @@ -6,8 +6,9 @@ then the expansion entry's reference): the block list and the token selection exactly. Columns past each row's visible count carry garbage (huge values, NaN) that must never be read. Beyond the cells: bitwise the composition it replaces (`qsa_topk_select`, then -`qsa_canonical_selection`, then `qsa_expand_block_indices`) on tie-free and tie-heavy rows; -a stale-read check under graph replay; one compiled kernel for every row count. +`qsa_canonical_selection`, then `qsa_expand_block_indices`, each certified on its own) on +one representative decode batch and one context chunk; a stale-read check under graph +replay; one compiled kernel for every row count. """ import zlib @@ -49,14 +50,6 @@ def _build(spec, seed): scores = torch.round(scores * 2.0) / 2.0 # few distinct values: ties at the threshold if spec.get("narrow"): scores = 1.0 + 0.05 * torch.rand(rows, W, generator=g, device=DEV) - if spec.get("const"): - scores = torch.full((rows, W), 0.75, device=DEV) - if spec.get("tied"): - # the 400 best distinct, then `tied` equal ones straddling the k-th place - scores = torch.rand(rows, W, generator=g, device=DEV) - best = torch.randperm(W // 4, generator=g, device=DEV) - scores[:, best[:400]] = 2.0 + torch.rand(rows, 400, generator=g, device=DEV) - scores[:, best[400 : 400 + spec["tied"]]] = 1.5 ends = ((pos + 1) // CR).to(torch.int32) for r in range(rows): # garbage past the visible blocks n = int(ends[r]) @@ -79,29 +72,21 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by assert torch.equal(sel, want_sel.to(torch.int32)), "token selection" -def _rows(kind: str, rows: int, seed: int) -> torch.Tensor: - """Score rows the canonical rewrite exists for: tie-free, relu-sparse (fewer positive - scores than the budget, so the zero class straddles it), two-valued and all-equal.""" +def _relu_sparse_rows(rows: int, seed: int) -> torch.Tensor: + """Score rows the canonical rewrite exists for: relu-sparse, so fewer positive scores + than the budget and the zero class (negative zeros among it) straddles it.""" g = torch.Generator(device=DEV).manual_seed(seed) - if kind == "tie_free": - return torch.rand(rows, W, generator=g, device=DEV) + 0.5 - if kind == "relu_sparse": - scores = torch.relu(torch.randn(rows, W, generator=g, device=DEV) - 2.2) - scores[:, ::7] = -0.0 # negative zero ties with zero - return scores - if kind == "two_valued": - return torch.where(torch.rand(rows, W, generator=g, device=DEV) < 0.3, 2.0, 1.0) - assert kind == "all_equal" - return torch.full((rows, W), 0.75, device=DEV) - - -@pytest.mark.parametrize("kind", ["tie_free", "relu_sparse", "two_valued", "all_equal"]) -def test_bitwise_the_canonical_composition(kind) -> None: + scores = torch.relu(torch.randn(rows, W, generator=g, device=DEV) - 2.2) + scores[:, ::7] = -0.0 # negative zero ties with zero + return scores + + +def test_bitwise_the_canonical_composition() -> None: """Windows across the budget (below it, at it, past it, the whole table); the block list equals `qsa_topk_select` + `qsa_canonical_selection` and the token selection that list's expansion, bitwise.""" rows = 64 - scores = _rows(kind, rows, seed=zlib.crc32(kind.encode()) % 2**31) + scores = _relu_sparse_rows(rows, seed=zlib.crc32(b"relu_sparse") % 2**31) g = torch.Generator(device=DEV).manual_seed(5) pos = torch.randint(0, CR * W, (rows,), generator=g, device=DEV, dtype=torch.int64) pos[:6] = torch.tensor([0, 4 * K - 1, 4 * K, 4 * K + 3, 8700, CR * W - 1], device=DEV) @@ -118,13 +103,12 @@ def test_bitwise_the_canonical_composition(kind) -> None: assert torch.equal(sel, want_sel.to(torch.int32)), "token selection" -@pytest.mark.parametrize("kind", ["tie_free", "relu_sparse"]) -def test_a_context_chunk_into_a_row_range_of_the_selection(kind) -> None: +def test_a_context_chunk_into_a_row_range_of_the_selection() -> None: """A context chunk's 2048 consecutive rows (one request, positions 6144..8191), written through `out` into rows [64, 2112) of a step-sized selection: the block list and those rows equal the canonical composition bitwise; the rows around them are untouched.""" rows = 2048 - scores = _rows(kind, rows, seed=31) + scores = _relu_sparse_rows(rows, seed=31) pos = torch.arange(6144, 6144 + rows, device=DEV, dtype=torch.int64) ends = ((pos + 1) // CR).to(torch.int32) seqlen = torch.full((rows,), 6144 + rows, dtype=torch.int32, device=DEV) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_select.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_select.py index b687789a1663..00cdbf65da13 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_select.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_select.py @@ -67,29 +67,18 @@ def _logits(kind: str, rows: int, cols: int, gen: torch.Generator) -> torch.Tens def _build(spec: dict, seed: int): gen = torch.Generator(device=DEV).manual_seed(seed) rows, cols, index_topk = spec["rows"], spec["cols"], spec["index_topk"] - if spec.get("strided"): - backing = _logits(spec["logits"], rows, 2 * cols, gen).view(rows, cols, 2) - logits = backing[:, :, 0] - assert logits.stride() == (2 * cols, 2) - else: - logits = _logits(spec["logits"], rows, cols, gen) - - kind = spec["windows"][0] - if kind in ("positions", "chunk"): - if kind == "positions": - positions = torch.tensor(spec["windows"][1], device=DEV) - else: - positions = spec["windows"][1] + torch.arange(rows, device=DEV) - assert positions.numel() == rows - # QSA's per-row causal bound: complete compressed blocks at the position. - row_ends = ((positions + 1) // COMPRESS_RATIO).clamp_max(cols).to(torch.int32) - row_starts = torch.zeros(rows, dtype=torch.int32, device=DEV) - elif kind == "explicit": - row_starts = torch.tensor(spec["windows"][1], dtype=torch.int32, device=DEV) - row_ends = torch.tensor(spec["windows"][2], dtype=torch.int32, device=DEV) + logits = _logits(spec["logits"], rows, cols, gen) + + kind, start = spec["windows"] + if kind == "positions": + positions = torch.tensor(start, device=DEV) else: - row_starts = torch.zeros(rows, dtype=torch.int32, device=DEV) - row_ends = torch.randint(0, cols + 1, (rows,), generator=gen, device=DEV).to(torch.int32) + assert kind == "chunk", kind + positions = start + torch.arange(rows, device=DEV) + assert positions.numel() == rows + # QSA's per-row causal bound: complete compressed blocks at the position. + row_ends = ((positions + 1) // COMPRESS_RATIO).clamp_max(cols).to(torch.int32) + row_starts = torch.zeros(rows, dtype=torch.int32, device=DEV) if spec["tail"] == "nan": columns = torch.arange(cols, device=DEV)[None, :] diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qwen4_exp_qk_norm_rope.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qwen4_exp_qk_norm_rope.py index 2cf8984e87c5..739b3b0cc340 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qwen4_exp_qk_norm_rope.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qwen4_exp_qk_norm_rope.py @@ -162,8 +162,7 @@ def _current_pair(proj, scratch, kwargs, index_k_weight): return proj, scratch[:, : spec["index_heads"]] -@pytest.mark.parametrize("tokens", [1, 19, 128]) -@pytest.mark.parametrize("rows", ["distinct", "text"]) +@pytest.mark.parametrize("tokens, rows", [(1, "distinct"), (128, "text")]) def test_matches_the_cuda_pair_it_replaces(tokens: int, rows: str) -> None: """Same wiring as the two CUDA calls: mRoPE sections (11, 10) on the main plane, coordinate 0 alone on the index plane, q/k weight assignment. Both @@ -290,8 +289,7 @@ def _bitwise(a, b) -> bool: return torch.equal(a.view(torch.int16), b.view(torch.int16)) -@pytest.mark.parametrize("tokens", [1024, 2048, 8193, 16384]) -@pytest.mark.parametrize("rows", ["distinct", "text"]) +@pytest.mark.parametrize("tokens, rows", [(1024, "distinct"), (8193, "text"), (16384, "distinct")]) def test_cute_kernel_is_the_triton_kernel_bitwise(tokens: int, rows: str) -> None: """From `_CUTE_MIN_TOKENS` tokens the op runs the CuTe kernel: every output bit (both planes, the untouched columns and scratch heads included) is the Triton kernel's.""" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/embedding/test_modeling_v2_ngram_hash.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/embedding/test_modeling_v2_ngram_hash.py index e88fb04f3fb5..91871a7408c3 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/embedding/test_modeling_v2_ngram_hash.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/embedding/test_modeling_v2_ngram_hash.py @@ -5,7 +5,7 @@ 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. +the tail slice and chunked prefill, and the four guards. """ import zlib @@ -81,7 +81,7 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res eos_positions=spec.get("eos_positions", ()), lead_eos=spec.get("lead_eos", True), eos_density=spec.get("eos_density", 0.0), - ).to(spec.get("ids_dtype", torch.int64)) + ) before = ids.clone() args = (ids, *_constants(), EOS, spec["num_query_tokens"]) out = op(*args) @@ -176,13 +176,6 @@ def test_eos_token_itself_ends_its_segment() -> None: 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)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_bf16_gemv_pdl.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_bf16_gemv_pdl.py index 0146f9a12052..bc27c7bfed48 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_bf16_gemv_pdl.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_bf16_gemv_pdl.py @@ -58,12 +58,12 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 _compare(out, op.reference(mat_a, mat_b, s["pdl"])) -@pytest.mark.parametrize("pdl", [False, True]) -def test_dependent_chain_under_graph_replay(pdl) -> None: +def test_dependent_chain_under_graph_replay() -> None: """Decode runs this under CUDA graphs, back to back with other kernels. Each link of the chain reads the output of the one before it: a GEMV that read its activation before `griddepcontrol.wait` would see the previous replay's values, and the comparison below would fail on the second replay.""" + pdl = True dims = [2560, 6144, 2560, 640, 2560] # K of each link; N is the next link's K weights = [_build(1, n, k, seed=i)[1] for i, (k, n) in enumerate(zip(dims, dims[1:] + [16480]))] x = torch.randn(1, dims[0], device="cuda").to(torch.bfloat16) @@ -97,18 +97,18 @@ def chain(): h = out -@pytest.mark.parametrize("m,n,k", [(1, 16480, 2560), (2, 2560, 6144), (1, 2560, 640)]) +@pytest.mark.parametrize("m,n,k", [(1, 16480, 2560), (1, 2560, 6144)]) def test_pdl_on_and_off_agree_bitwise(m, n, k) -> None: """PDL changes when the kernel may start, not what it computes.""" mat_a, mat_b = _build(m, n, k, seed=11) assert torch.equal(launch(mat_a, mat_b, True), launch(mat_a, mat_b, False)) -@pytest.mark.parametrize("m,n,k", [(1, 16480, 2560), (2, 2560, 6144)]) -def test_band_discriminates_wrong_variants(m, n, k) -> None: +def test_band_discriminates_wrong_variants() -> None: """The band admits the kernel and rejects the near misses a wrong k-tile split or layout would produce: half of K missing, a one-column shift of the output, a 3% scale error.""" + m, n, k = 1, 16480, 2560 mat_a, mat_b = _build(m, n, k, seed=7) out = op(mat_a, mat_b) ref = op.reference(mat_a, mat_b) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_low_m_bf16_mm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_low_m_bf16_mm.py index f788ef2dffcb..303ea44b6ecd 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_low_m_bf16_mm.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_low_m_bf16_mm.py @@ -2,7 +2,8 @@ # SPDX-License-Identifier: Apache-2.0 """GPU test for the low_m_bf16_mm entry. -Every cell names the branch `plan` must pick, so a table edit that moves a +Every cell names the branch `plan` must pick, and a host-only check pins the +branch of the table's other shapes and M buckets, so a table edit that moves a shape to another kernel fails here before it can change a number silently. Beyond the cells: CUDA-graph capture and replay after warmup (the way the decode forward runs it), the wrong-variant discrimination of the band, and the @@ -60,11 +61,46 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 _compare(out, op.reference(mat_a, mat_b, s["pdl"])) -@pytest.mark.parametrize("pdl", [False, True]) -def test_graph_replay_after_warmup(pdl) -> None: - """Decode runs this under CUDA graphs: compile in warmup, capture, then replay - with new activations copied into the captured input. One call per branch, in - a chain, so a FlashInfer kernel follows cuBLAS and vice versa.""" +#: `(M, N, K, branch)` of the table's shapes and M buckets the cells do not drive: +#: what `plan` picks for each, held without a launch. +_PLAN_PINS = ( + (1, 10240, 2560, "direct"), + (1, 48, 2560, "direct"), + (1, 2560, 6144, "direct"), + (3, 12288, 2560, "splitk"), + (8, 12288, 2560, "splitk"), + (16, 12288, 2560, "cublas"), + (16, 512, 2560, "direct"), + (24, 512, 2560, "splitk"), + (16, 640, 2560, "splitk"), + (4, 1280, 2560, "direct"), + (8, 1296, 2560, "splitk"), + (1, 1, 2560, "direct"), + (1, 16480, 2560, "direct"), + (2, 14080, 2560, "direct"), + (3, 14080, 2560, "splitk"), + (8, 14080, 2560, "splitk"), + (16, 14080, 2560, "cublas"), + (1, 2560, 2560, "direct"), + (4, 2560, 2560, "direct"), + (8, 2560, 2560, "splitk"), + (1, 2560, 640, "cublas"), + (8192, 2560, 6144, "cublas"), + (1, 1024, 2048, "cublas"), +) + + +def test_plan_pins_the_branch_of_the_other_table_shapes() -> None: + picked = [(m, n, k, want, plan(m, n, k).kind) for m, n, k, want in _PLAN_PINS] + wrong = [p for p in picked if p[3] != p[4]] + assert not wrong, f"(M, N, K, pinned, picked): {wrong}" + + +def test_graph_replay_after_warmup() -> None: + """Decode runs this under CUDA graphs, with PDL: compile in warmup, capture, + then replay with new activations copied into the captured input. One call per + branch, in a chain, so a FlashInfer kernel follows cuBLAS and vice versa.""" + pdl = True shapes = [(1, 10240, 2560), (8, 2560, 6144), (1, 2560, 640), (8, 12288, 2560)] assert [plan(*s).kind for s in shapes] == ["direct", "splitk", "cublas", "splitk"] ins = [_build(m, n, k, seed=i) for i, (m, n, k) in enumerate(shapes)] @@ -88,11 +124,11 @@ def test_graph_replay_after_warmup(pdl) -> None: _compare(out, op.reference(a, b)) -@pytest.mark.parametrize("m,n,k", [(1, 10240, 2560), (8, 2560, 6144)]) -def test_band_discriminates_wrong_variants(m, n, k) -> None: +def test_band_discriminates_wrong_variants() -> None: """The band admits the kernel and rejects the near misses a wrong layout or a dropped K-split would produce: half of K missing, a one-column shift of the output, a 3% scale error.""" + m, n, k = 8, 2560, 6144 mat_a, mat_b = _build(m, n, k, seed=7) out = op(mat_a, mat_b) ref = op.reference(mat_a, mat_b) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py index 9a7ffbe1e12f..2ca906f8b1e1 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py @@ -100,14 +100,10 @@ def normal(*shape, scale): return torch.randn(*shape, generator=g, device=DEV) * scale hyper_input = normal(rows, width, scale=3.0).to(dtype) - if spec.get("layout") == "row_slice": - buf = normal(2 * rows, width, scale=1.0).to(dtype) - buf[::2] = hyper_input - hyper_input = buf[::2] gamma = normal(width, scale=GAMMA_RMS) gamma[0], gamma[width // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX gamma = gamma.to(dtype) - n_rows = spec.get("n_rows", -(-n_cols // 128) * 128) + n_rows = -(-n_cols // 128) * 128 weight = torch.zeros(n_rows, width, device=DEV) weight[:n_cols] = normal(n_cols, width, scale=DOWN_RMS) weight[n_cols - hc_count : n_cols] *= INJECT_RMS / DOWN_RMS @@ -119,8 +115,6 @@ def normal(*shape, scale): if spec.get("moe", False): shared = normal(rows, hidden, scale=2.0).to(dtype) gate_logits = normal(rows, 1, scale=2.0).to(dtype) - if spec.get("gate") == "flat": - gate_logits = gate_logits.reshape(rows) top_k = spec.get("fold", 0) if not top_k: return hyper_input, block_out, inject, gamma, first, hc_count, n_cols, shared, gate_logits diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_cute.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_cute.py deleted file mode 100644 index 93e21cd02dca..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_cute.py +++ /dev/null @@ -1,118 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the qwen4_exp_hc_decode_mix_cute entry (the plain build of the CuTe mix -kernel): cells against the entry's float64 reference in its band (`mixed` and -`inject_out` in `rtol = eps`, `atol = eps * max|ref|`), a stale-read check under -CUDA-graph replay with PDL, 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.hc.qwen4_exp_hc_decode_mix_fp4_cute import ( # noqa: E501 - qwen4_exp_hc_decode_mix_cute as op, -) - -__extra_import_path__ = ["../../../.."] -from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name - -assert torch.cuda.is_available(), "qwen4_exp_hc_decode_mix_cute requires a CUDA device" - -DEV = "cuda" -HC_COUNT, HIDDEN, RANK, EPS = 4, 2560, 320, 1e-6 -UP_RMS, GAMMA_RMS, GAMMA_ABSMAX = 0.0681, 0.477, 6.625 - - -def _band(got, want) -> float: - eps = torch.finfo(want.dtype).eps - w = want.double() - return ((got.double() - w).abs() / (eps * w.abs().max().item() + eps * w.abs())).max().item() - - -def _compare(out, ref) -> None: - (mixed, inj), (ref_mixed, ref_inj) = out, ref - assert (inj is None) == (ref_inj is None) - for got, want in [(mixed, ref_mixed)] + ([] if inj is None else [(inj, ref_inj)]): - assert got.dtype == want.dtype and got.shape == want.shape and got.is_contiguous() - assert _band(got, want) <= 1.0, f"{_band(got, want):.4g}x the band" - - -def _build(spec, seed): - rows, hc_count, hidden, rank, dtype = 1, HC_COUNT, HIDDEN, RANK, torch.bfloat16 - with_inject = spec.get("with_inject", True) - width = hc_count * hidden - g = torch.Generator(device=DEV).manual_seed(seed) - - def normal(*shape, scale): - return torch.randn(*shape, generator=g, device=DEV) * scale - - hc_out = normal(rows, width, scale=3.0).to(dtype) - if spec.get("layout") == "row_slice": - buf = normal(2 * rows, width, scale=1.0).to(dtype) - buf[::2] = hc_out - hc_out = buf[::2] - ss = hc_out.double().unflatten(-1, (hc_count, hidden)).pow(2).sum(-1).t().float().contiguous() - gamma = normal(width, scale=GAMMA_RMS) - gamma[0], gamma[width // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX - gamma = gamma.to(dtype) - n_pad = -(-(rank + hc_count) // 4) * 4 - part = normal(hc_count, rows, n_pad + 4, scale=60.0)[:, :, :n_pad] - up = normal(width, rank, scale=UP_RMS).to(dtype).t() - return [part, ss, up, hc_out, gamma, EPS, hc_count, with_inject] - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name - certifying(op) - args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) - out = op(*args) - assert out[0].shape == (1, HIDDEN) - _compare(out, op.reference(*args)) - - -def test_deterministic() -> None: - args = _build(dict(), seed=3) - for a, b in zip(op(*args), op(*args)): - assert torch.equal(a, b) - - -def test_reads_after_the_wait_under_graph_replay() -> None: - """Producer kernels rewrite `part`, `ss` and `hc_out` inside the graph right before - the op; with PDL a read before `griddepcontrol.wait` would see the previous replay's.""" - args = _build(dict(), seed=8) - targets = (args[0], args[1], args[3]) - sources = [t.clone() for t in targets] - - def step(): - for dst, src in zip(targets, sources): - dst.copy_(src * 1.0) - return op(*args) - - warmup = torch.cuda.Stream() - warmup.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(warmup): - for _ in range(2): - step() - torch.cuda.current_stream().wait_stream(warmup) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - graph_out = step() - for trial in range(3): - fresh = _build(dict(), seed=100 + trial) - for src, new in zip(sources, (fresh[0], fresh[1], fresh[3])): - src.copy_(new) - graph.replay() - torch.cuda.synchronize() - for a, b in zip(graph_out, op(*args)): - assert torch.equal(a, b), f"replay {trial}" - - -def test_guards() -> None: - args = _build(dict(), seed=6) - two = args[3].expand(2, -1).contiguous() - with validating(op), pytest.raises(AssertionError): - op(args[0].expand(-1, 2, -1), args[1].expand(-1, 2), args[2], two, *args[4:]) - contiguous_up = args[2].contiguous() - with validating(op), pytest.raises(AssertionError): - op(args[0], args[1], contiguous_up, *args[3:]) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py index 761212e90c47..b92d780b5cde 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py @@ -1,10 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the qwen4_exp_hc_decode_mix_fp4_cute entry: cells against the entry's -float64 reference in its band (`mixed` and `inject_out` in `rtol = eps`, -`atol = eps * max|ref|`; the NVFP4 exact on the returned `mixed` outside the -reciprocal tie window), a stale-read check under CUDA-graph replay with PDL, the -guards.""" +"""GPU test for the two entries of qwen4_exp_hc_decode_mix_fp4_cute: the NVFP4 build +(`qwen4_exp_hc_decode_mix_fp4_cute`) and the plain build (`qwen4_exp_hc_decode_mix_cute`) +of the CuTe mix kernel. Each entry's cells against its float64 reference in its band +(`mixed` and `inject_out` in `rtol = eps`, `atol = eps * max|ref|`; for the NVFP4 build, +the NVFP4 exact on the returned `mixed` outside the reciprocal tie window), a +determinism check, a stale-read check under CUDA-graph replay with PDL for each build, +the guards.""" import zlib @@ -16,7 +18,8 @@ fp4_quantize, ) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_decode_mix_fp4_cute import ( # noqa: E501 - qwen4_exp_hc_decode_mix_fp4_cute as op, + qwen4_exp_hc_decode_mix_cute, + qwen4_exp_hc_decode_mix_fp4_cute, ) __extra_import_path__ = ["../../../.."] @@ -28,6 +31,9 @@ HC_COUNT, HIDDEN, RANK, EPS = 4, 2560, 320, 1e-6 UP_RMS, GAMMA_RMS, GAMMA_ABSMAX = 0.0681, 0.477, 6.625 +ENTRIES = {"fp4": qwen4_exp_hc_decode_mix_fp4_cute, "plain": qwen4_exp_hc_decode_mix_cute} +CASES = [(name, cell) for name, entry in ENTRIES.items() for cell in entry.CELLS] + def nvfp4_reference(x, gs): """`fp4_quantize`'s native-torch NVFP4 of `x` `[T, K]` with linear scales: @@ -70,7 +76,8 @@ def _check_quant(mixed, xq, xsf, gs) -> None: assert not stray.any(), f"{stray.sum().item()} e2m1 codes differ outside the tie window" -def _build(spec, seed): +def _build(spec, seed, fp4=True): + """The entry's arguments; `fp4` appends the NVFP4 build's global scale.""" rows, hc_count, hidden, rank, dtype = 1, HC_COUNT, HIDDEN, RANK, torch.bfloat16 with_inject = spec.get("with_inject", True) width = hc_count * hidden @@ -91,31 +98,42 @@ def normal(*shape, scale): n_pad = -(-(rank + hc_count) // 4) * 4 part = normal(hc_count, rows, n_pad + 4, scale=60.0)[:, :, :n_pad] up = normal(width, rank, scale=UP_RMS).to(dtype).t() - # A calibrated activation global scale: 448 * 6 / amax of the mix (~|8|). - gs = torch.tensor([448.0 * 6.0 / 8.0], device=DEV, dtype=torch.float32) - return [part, ss, up, hc_out, gamma, EPS, hc_count, with_inject, gs] + args = [part, ss, up, hc_out, gamma, EPS, hc_count, with_inject] + if fp4: + # A calibrated activation global scale: 448 * 6 / amax of the mix (~|8|). + args.append(torch.tensor([448.0 * 6.0 / 8.0], device=DEV, dtype=torch.float32)) + return args -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name +@pytest.mark.parametrize( + "name,cell", CASES, ids=[f"{name}-{cell.why[:44]}" for name, cell in CASES] +) +def test_certified_cells(name, cell, certifying) -> None: # noqa: F811 — fixture by name + op = ENTRIES[name] certifying(op) - args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) + fp4 = name == "fp4" + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31), fp4=fp4) out = op(*args) assert out[0].shape == (1, HIDDEN) _compare(out, op.reference(*args)) - _check_quant(out[0], out[2], out[3], args[8]) + if fp4: + _check_quant(out[0], out[2], out[3], args[8]) def test_deterministic() -> None: + op = qwen4_exp_hc_decode_mix_fp4_cute args = _build(dict(), seed=3) for a, b in zip(op(*args), op(*args)): assert torch.equal(a, b) -def test_reads_after_the_wait_under_graph_replay() -> None: +@pytest.mark.parametrize("name", list(ENTRIES)) +def test_reads_after_the_wait_under_graph_replay(name) -> None: """Producer kernels rewrite `part`, `ss` and `hc_out` inside the graph right before - the op; with PDL a read before `griddepcontrol.wait` would see the previous replay's.""" - args = _build(dict(), seed=8) + the op; with PDL a read before `griddepcontrol.wait` would see the previous replay's. + Each build is its own compiled variant, so each is replayed.""" + op, fp4 = ENTRIES[name], name == "fp4" + args = _build(dict(), seed=8, fp4=fp4) targets = (args[0], args[1], args[3]) sources = [t.clone() for t in targets] @@ -134,7 +152,7 @@ def step(): with torch.cuda.graph(graph): graph_out = step() for trial in range(3): - fresh = _build(dict(), seed=100 + trial) + fresh = _build(dict(), seed=100 + trial, fp4=fp4) for src, new in zip(sources, (fresh[0], fresh[1], fresh[3])): src.copy_(new) graph.replay() @@ -144,6 +162,8 @@ def step(): def test_guards() -> None: + """Both entries share the shape checks; the NVFP4 build's `is_valid` runs them.""" + op = qwen4_exp_hc_decode_mix_fp4_cute args = _build(dict(), seed=6) two = args[3].expand(2, -1).contiguous() with validating(op), pytest.raises(AssertionError): diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first.py index 8ba9eb3d73e6..2bceb75ebecb 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first.py @@ -47,7 +47,7 @@ def _compare(out, ref) -> None: def _build(spec, seed): - hc, hidden = spec.get("hc_count", 4), spec.get("hidden", 2560) + hc, hidden = 4, 2560 rows, width, inject = spec["rows"], hc * hidden, spec.get("inject", True) g = torch.Generator(device=DEV).manual_seed(seed) # This checkpoint's layer statistics: down ~N(0, 0.0296), block_inject ~N(0, 0.0107), @@ -55,9 +55,7 @@ def _build(spec, seed): down = torch.randn(RANK, width, generator=g, device=DEV) * 0.0296 parts = [down] + ([torch.randn(hc, width, generator=g, device=DEV) * 0.0107] if inject else []) w = torch.cat(parts) - n_rows = spec.get("n_rows", 384) - if n_rows > w.shape[0]: - w = torch.cat([w, torch.zeros(n_rows - w.shape[0], width, device=DEV)]) + w = torch.cat([w, torch.zeros(384 - w.shape[0], width, device=DEV)]) gamma = torch.randn(width, generator=g, device=DEV) * 0.477 x = torch.randn(rows, hc, hidden, generator=g, device=DEV) x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) @@ -131,13 +129,13 @@ def test_cuda_graph_capture_and_replay() -> None: assert torch.equal(g_low, low) and torch.equal(g_inject, inject) -@pytest.mark.parametrize("rows", [2, 33, 128]) -@pytest.mark.parametrize("with_inject", [True, False]) -def test_mixer_matches_the_gated_residual_reference(rows, with_inject) -> None: +def test_mixer_matches_the_gated_residual_reference() -> None: """The decode mixer as `DecodeTarget._hc_mixer` runs it from 2 rows -- this op, then `qwen4_exp_hc_up_gate_mean` -- against `qwen4_exp_gated_residual.reference` (float64 from `normed`) in the band that entry certifies its own op in, beside - the cuBLAS + `qwen4_exp_hc_lowrank_act` mixer it replaces.""" + the cuBLAS + `qwen4_exp_hc_lowrank_act` mixer it replaces. One representative + case (128 rows, with injection): each component certifies its own cells.""" + rows, with_inject = 128, True from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501 gate_ratio, diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_inject_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_inject_norm.py index aa6de075afca..0341002b3ddf 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_inject_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_inject_norm.py @@ -124,17 +124,10 @@ def _compare_moe(args, out, ref) -> None: def _build(spec, seed): - rows = spec["rows"] - hc_count = spec.get("hc_count", HC_COUNT) - hidden = spec.get("hidden", HIDDEN) - dtype = spec.get("dtype", torch.bfloat16) + rows, hc_count, hidden, dtype = spec["rows"], HC_COUNT, HIDDEN, torch.bfloat16 width = hc_count * hidden g = torch.Generator(device=DEV).manual_seed(seed) hyper_input = (torch.randn(rows, width, generator=g, device=DEV) * 3.0).to(dtype) - if spec.get("layout") == "row_slice": - buf = torch.randn(2 * rows, width, generator=g, device=DEV).to(dtype) - buf[::2] = hyper_input - hyper_input = buf[::2] gamma = torch.randn(width, generator=g, device=DEV) * GAMMA_RMS gamma[0], gamma[width // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX gamma = gamma.to(dtype) @@ -147,8 +140,6 @@ def _build(spec, seed): if spec.get("moe", False): shared = (torch.randn(rows, hidden, generator=g, device=DEV) * 2.0).to(dtype) gate_logits = (torch.randn(rows, 1, generator=g, device=DEV) * 2.0).to(dtype) - if spec.get("gate") == "flat": - gate_logits = gate_logits.reshape(rows) top_k = spec.get("fold", 0) if not top_k: return hyper_input, block_out, inject, gamma, EPS, hc_count, shared, gate_logits @@ -204,10 +195,11 @@ def _old_composition(hyper_input, block_out, inject, gamma, eps, hc_count): return hc, normed -@pytest.mark.parametrize("rows", [1, 128, 8192]) -def test_matches_the_composition_it_replaces(rows) -> None: +def test_matches_the_composition_it_replaces() -> None: """Against post_inject + four gemma norms + cat: the residual within one - ulp (fma contraction only), the normed stream within one ulp.""" + ulp (fma contraction only), the normed stream within one ulp. One + representative prefill chunk: each component certifies its own cells.""" + rows = 8192 args = _build(dict(rows=rows), seed=rows)[:6] hc, normed = op(*args) old_hc, old_normed = _old_composition(*args) @@ -281,11 +273,12 @@ def test_cuda_graph_capture_and_replay() -> None: assert torch.equal(graph_normed, eager_normed) -@pytest.mark.parametrize("rows", [2, 128, 2048]) +@pytest.mark.parametrize("rows", [128, 2048]) def test_fold_matches_the_runner_finalize(rows) -> None: """With the fold, bitwise this op on the runner's finalized rows: `finalize` (the permuted MoE entry's emulation of the runner's finalize, bitwise its - output) and then the combine without the fold.""" + output) and then the combine without the fold. 128 rows run the CuTe kernel, + 2048 the Triton kernel's two-warp form.""" args = _build(dict(rows=rows, moe=True, fold=10), seed=300 + rows) hyper_input, rows_, inject, gamma, eps, c, shared, gate, topk_weights, topk_rows = args hc, normed = op(*args) @@ -295,10 +288,11 @@ def test_fold_matches_the_runner_finalize(rows) -> None: assert torch.equal(normed, old_normed) -@pytest.mark.parametrize("rows", [1, 128, 8192]) -def test_moe_combine_matches_the_composition_it_replaces(rows) -> None: +def test_moe_combine_matches_the_composition_it_replaces() -> None: """Against `fused_sigmoid_gate_mul_add` (in place on the routed output, as - the forward called it) and then this op without the combine: bitwise.""" + the forward called it) and then this op without the combine: bitwise, at one + representative decode batch.""" + rows = 128 args = _build(dict(rows=rows, moe=True), seed=100 + rows) hyper_input, block_out, inject, gamma, eps, c, shared, gate = args hc, normed = op(*args) @@ -382,13 +376,8 @@ def test_moe_cuda_graph_capture_and_replay() -> None: @pytest.mark.parametrize( "spec", [ - dict(rows=1), dict(rows=2, drop=True), - dict(rows=8), - dict(rows=33, drop=True), dict(rows=128), - dict(rows=17, gate="flat", layout="row_slice"), - dict(rows=300, drop=True), dict(rows=511), ], ids=lambda s: "-".join(f"{k}{v}" for k, v in s.items()), @@ -396,7 +385,8 @@ def test_moe_cuda_graph_capture_and_replay() -> None: def test_fold_cute_is_the_triton_kernel_bitwise(spec, monkeypatch) -> None: """Below `_FEW_WARPS_MIN_ROWS` rows the fold runs the CuTe DSL kernel; its `hc_out` and `normed` are bitwise the Triton kernel's (run by refusing the - CuTe guard).""" + CuTe guard): the smallest decode batch with dropped slots, the largest decode + batch, and the last row count before the switch.""" args = _build(dict(spec, moe=True, fold=10), seed=500 + spec["rows"]) calls = [] cute_fold = M._fold_cute diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py index 3521a59f384a..cea97275bf3a 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py @@ -99,13 +99,13 @@ def test_cuda_graph_capture_and_replay() -> None: assert torch.equal(g_mixed, op(fresh[0], up, fresh[2], c)) -@pytest.mark.parametrize("rows", [1, 33, 128, 2048, 8192, 16384]) -@pytest.mark.parametrize("with_inject", [True, False]) -def test_mixer_matches_the_gated_residual_reference(rows, with_inject) -> None: +def test_mixer_matches_the_gated_residual_reference() -> None: """The mixer as the forward runs it -- `cublas_mm` over the zero-padded `[down; inject]` (fp32 out), `qwen4_exp_hc_lowrank_act`, this op -- against `qwen4_exp_gated_residual.reference` (float64 from `normed`), in the band - that entry certifies its own op in.""" + that entry certifies its own op in. One prefill-sized case: each component + is certified over its own cells, and the injection columns never reach + `mixed`.""" from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501 qwen4_exp_gated_residual, @@ -114,7 +114,7 @@ def test_mixer_matches_the_gated_residual_reference(rows, with_inject) -> None: qwen4_exp_hc_lowrank_act, ) - hc, hidden, rank, eps = 4, 2560, 320, 1e-6 + rows, hc, hidden, rank, eps = 2048, 4, 2560, 320, 1e-6 width = hc * hidden g = torch.Generator(device=DEV).manual_seed(rows) @@ -128,16 +128,14 @@ def normal(*shape, scale): x = torch.randn(rows, hc, hidden, generator=g, device=DEV) x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) normed = (x.flatten(-2) * (1.0 + gamma.float())).to(torch.bfloat16) - first = torch.cat([down, inject_w], 0) if with_inject else down + first = torch.cat([down, inject_w], 0) first = torch.cat( [first, torch.zeros(384 - first.shape[0], width, device=DEV, dtype=first.dtype)] ) low, _ = qwen4_exp_hc_lowrank_act( - cublas_mm(normed, first.t(), None, torch.float32), rank, hc, with_inject, torch.bfloat16 + cublas_mm(normed, first.t(), None, torch.float32), rank, hc, True, torch.bfloat16 ) mixed = op(low, up.t(), normed, hc) - want = qwen4_exp_gated_residual.reference( - normed, down, up, inject_w if with_inject else None, hc - ) + want = qwen4_exp_gated_residual.reference(normed, down, up, inject_w, hc) _compare(mixed, want[0]) print(f"\n rows {rows}: fused {gate_ratio(mixed, want[0]):.3f}x band") diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4.py index aa4b31487bdf..17cb0509f5ff 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4.py @@ -105,16 +105,15 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by _compare(out, op.reference(*args), args[4]) -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_bitwise_the_composition_it_replaces(cell) -> None: +@pytest.mark.parametrize("rows", [2, 16390]) +def test_bitwise_the_composition_it_replaces(rows) -> None: """`qwen4_exp_hc_up_gate_mean` then `torch.ops.trtllm.fp4_quantize` (linear - scales): every output bitwise equal, ties included.""" - low, up, normed, c, gs = _build(cell.spec, seed=_seed(cell)) + scales): every output bitwise equal, ties included. A one-tile decode batch + and a many-tile one with a ragged tail; the cells certify the rest.""" + low, up, normed, c, gs = _build(dict(rows=rows), seed=rows) mixed, xq, xsf = op(low, up, normed, c, gs) old = qwen4_exp_hc_up_gate_mean(low, up, normed, c) assert torch.equal(mixed, old) - if old.shape[0] == 0: - return old_q, old_sf = torch.ops.trtllm.fp4_quantize(old, gs, 16, False, False) assert torch.equal(xq, old_q.view(torch.uint8)), ( f"{(xq != old_q.view(torch.uint8)).sum().item()} bytes" @@ -122,9 +121,10 @@ def test_bitwise_the_composition_it_replaces(cell) -> None: assert torch.equal(xsf, old_sf.view(torch.uint8).reshape(-1)) -@pytest.mark.parametrize("scale", [1e-6, 1e-3, 1.0, 448.0 * 6.0 / 24.0, 1e4]) +@pytest.mark.parametrize("scale", [1e-6, 448.0 * 6.0 / 24.0, 1e4]) def test_bitwise_across_global_scales(scale) -> None: - """Tiny scales (zero / subnormal e4m3 bytes) and saturating ones.""" + """A tiny scale (zero / subnormal e4m3 bytes), the calibrated one and a + saturating one.""" low, up, normed, c, _ = _build(dict(rows=128), seed=11) h = normed.shape[1] // c for s in range(c): # row 0's first 64 mix columns are exactly zero: amax 0 diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py index df6d50d69c22..ea92991e4625 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py @@ -7,7 +7,8 @@ it, the sort as moe_sort returns it -- and holds the claims no cell carries: per-pair invariance across tiles, tactics, row order and prefixes; CUDA-graph replay at the captured row counts; the regions the op leaves unwritten; that -the gate rejects every wrong layout; the guards and the loud rejections. +the gate rejects a wrong layout of each operand; the guards and the loud +rejections. """ import functools @@ -59,10 +60,11 @@ FLIP_BUDGET = 1e-5 #: Row counts of a decode capture (requests x MTP verify width 1..4, up to 32 -#: requests), and the eager prefix lengths: both sides of the power-of-two edges up -#: to the 8192-row ceiling. -_CAPTURED_ROWS = (1, 2, 3, 4, 5, 8, 32, 96, 128) -_PREFIXES = (1, 2, 3, 4, 5, 8, 31, 32, 33, 127, 128, 129, 512, 1023, 1024, 1025, 4095, 4096, 8191) +#: requests): both ends of the sort's block kernel, its first cluster-kernel T and +#: the largest capture. The eager prefix lengths straddle the captured ceiling and +#: reach the 8192-row one. +_CAPTURED_ROWS = (1, 4, 5, 128) +_PREFIXES = (1, 4, 5, 127, 128, 129, 4096, 8191) # ── the gate ────────────────────────────────────────────────────────────── @@ -208,9 +210,6 @@ def _row_perm(n: int, layout: str) -> torch.Tensor: return idx.view(2, n // (2 * GATE_UP_GROUP), GATE_UP_GROUP).transpose(0, 1).reshape(-1) if layout == "rowwise": # trtllm-gen's interleave, without its 32-row shuffle return idx.view(2, n // 2).transpose(0, 1).reshape(-1) - if layout == "gate_first": - half = idx.view(2, n // 2).flip(0).reshape(-1) - return half[_row_perm(n, "group64")] raise ValueError(layout) @@ -220,15 +219,10 @@ def _weights(layout: str = "group64", num_experts: int = E): target's load-time transform produces -- or one of the wrong layouts the discrimination test feeds it.""" w, s, ws2 = _experts(num_experts) - perm = _row_perm( - 2 * INTER, "group64" if layout in ("linear_scales", "scales_unpermuted") else layout - ) + perm = _row_perm(2 * INTER, "group64" if layout == "scales_unpermuted" else layout) w_k = w[:, perm].contiguous() s_rows = s if layout == "scales_unpermuted" else s[:, perm] - if layout == "linear_scales": - s_k = s_rows.contiguous() - else: - s_k = torch.ops.trtllm.block_scale_interleave(s_rows.contiguous()).view(s.shape) + s_k = torch.ops.trtllm.block_scale_interleave(s_rows.contiguous()).view(s.shape) alpha = (ISC1 * ws2).contiguous() return w_k.view(torch.float4_e2m1fn_x2), s_k, alpha @@ -382,7 +376,7 @@ def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves th # ── invariance: a pair's output depends on the pair alone ───────────────── -@pytest.mark.parametrize("tokens", (1, 128, 8192)) +@pytest.mark.parametrize("tokens", (128, 8192)) def test_tiles_and_tactics_agree_per_pair(tokens, certifying): # noqa: F811 """Tile 128 / 256, MMA N 128 / 256 and the tuner's fallback all produce every pair's codes and block scales bit for bit the same, so a cell certifies the @@ -404,13 +398,12 @@ def test_tiles_and_tactics_agree_per_pair(tokens, certifying): # noqa: F811 _assert_pairs_equal(got, first, f"T={tokens} {key} vs (128, 128)") -@pytest.mark.parametrize("tile", (128, 256)) -def test_row_order_and_repeat_do_not_change_a_pair(tile, certifying): # noqa: F811 +def test_row_order_and_repeat_do_not_change_a_pair(certifying): # noqa: F811 """moe_sort orders an expert's rows by atomics (from T = 5). Shuffling every expert's run of `permuted_idx_to_expanded_idx`, and calling twice, leaves each pair's output bitwise unchanged: the FC1 result is reproducible run to run.""" certifying(op) - tokens = 4096 + tokens, tile = 4096, 128 args, sort = _problem(tokens, tile, seed=77) base = _pairs(op(**args), sort, tokens) _assert_pairs_equal(_pairs(op(**args), sort, tokens), base, "repeat call") @@ -435,12 +428,11 @@ def test_row_order_and_repeat_do_not_change_a_pair(tile, certifying): # noqa: F _assert_pairs_equal(_pairs(out_s, shuffled, tokens), base, "shuffled rows") -@pytest.mark.parametrize("tile", (128, 256)) -def test_eager_rows_prefix_bitwise(tile, certifying): # noqa: F811 +def test_eager_rows_prefix_bitwise(certifying): # noqa: F811 """The 8192-row call is gated against the reference, and every prefix call's pairs are bitwise the 8192-row call's.""" certifying(op) - tokens = 8192 + tokens, tile = 8192, 128 x, x_sf = _activation(tokens, seed=8) ids = _ids("router", tokens, tile, seed=9) sort = _sort(ids, tile) @@ -473,13 +465,13 @@ def _capture(fn): return graph, static -@pytest.mark.parametrize("tile", (128, 256)) -def test_captured_rows_replay_on_new_data(tile, certifying): # noqa: F811 +def test_captured_rows_replay_on_new_data(certifying): # noqa: F811 """moe_sort and this op are captured together, as the target's graph holds them; each graph is replayed on two fresh activations and id sets, and its written rows are gated against the reference and bitwise against an eager call on the - graph's own sort output.""" + graph's own sort output. At the target's tile, 128.""" certifying(op) + tile = 128 for tokens in _CAPTURED_ROWS: x, x_sf = _activation(tokens, seed=tokens) ids = _ids("router", tokens, tile, seed=500 + tokens) @@ -539,14 +531,14 @@ def step(ids=ids, scales=scales, base=base, state=state, tile=tile): torch.cuda.synchronize() -@pytest.mark.parametrize("tile", (128, 256)) @pytest.mark.parametrize("tokens", (3, 300)) -def test_unwritten_regions(tokens, tile): +def test_unwritten_regions(tokens): """A sentinel planted in a captured graph's outputs: which bytes replay writes. The rows the sort filled hold the result; the padding rows of a written tile are written with zero block scales and zero codes; the tiles past `num_non_exiting_tiles` are not written -- under a graph they keep whatever - the buffer held.""" + the buffer held. Tile 128.""" + tile = 128 args, sort = _problem(tokens, tile, seed=11) graph, static_out = _capture(lambda: op(**args)) c, c_sf = static_out @@ -579,21 +571,14 @@ def test_unwritten_regions(tokens, tile): @pytest.mark.parametrize( "layout", - ( - "rowwise", - "gate_first", - "linear_scales", - "scales_unpermuted", - "swizzled_input_scales", - "alpha_rolled", - ), + ("rowwise", "scales_unpermuted", "swizzled_input_scales", "alpha_rolled"), ) def test_wrong_layouts_are_rejected(layout): - """Each operand prepared the wrong way -- trtllm-gen's row interleave, gate - rows first, scales left linear, scales not given the weight's row - permutation, the activation's scales swizzled, alpha off by one expert -- - returns finite, plausibly scaled output that `_compare` must reject against - the reference of the correctly prepared operands.""" + """Each operand prepared the wrong way -- the weight in trtllm-gen's row + interleave, its scales not given the weight's row permutation, the + activation's scales swizzled, alpha off by one expert -- returns finite, + plausibly scaled output that `_compare` must reject against the reference of + the correctly prepared operands.""" tokens, tile = 128, 128 x, x_sf = _activation(tokens, seed=3) sort = _sort(_ids("router", tokens, tile, seed=4), tile) @@ -632,10 +617,12 @@ def _u8(t: torch.Tensor) -> torch.Tensor: return t.view(torch.uint8) if t.element_size() == 1 else t -@pytest.mark.parametrize("name", ("input", "input_scale", "weight", "weight_scale", "alpha")) +@pytest.mark.parametrize("name", ("input", "weight_scale")) def test_noncontiguous_operand_answers_wrongly_and_is_refused(name): """The runner hands data_ptr() to the kernel and never reads a stride: the - same values through a strided view come back finite and wrong.""" + same values through a strided view come back finite and wrong. One activation + and one weight operand; the guard checks all five of input, input_scale, + weight, weight_scale and alpha the same way.""" tokens, tile = 64, 128 args, _ = _problem(tokens, tile, seed=21) ref = op.reference(**args) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_grouped_gemm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_grouped_gemm.py index 00e319236451..f86e904bc5e5 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_grouped_gemm.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_grouped_gemm.py @@ -8,8 +8,8 @@ unwritten bytes past the live tiles), the sort as moe_sort returns it -- and holds the claims no cell carries: per-pair invariance across tiles, tactics, row order and prefixes; CUDA-graph replay of sort, FC1 and FC2 at the captured -row counts; the regions the op leaves unwritten; that the gate rejects every -wrong layout; the guards and the loud rejections. +row counts; the regions the op leaves unwritten; that the gate rejects a +wrong layout of each operand; the guards and the loud rejections. """ import functools @@ -55,10 +55,11 @@ FLIP_BUDGET = 1e-5 #: Row counts of a decode capture (requests x MTP verify width 1..4, up to 32 -#: requests), and the eager prefix lengths: both sides of the power-of-two edges up -#: to the 8192-row ceiling. -_CAPTURED_ROWS = (1, 2, 3, 4, 5, 8, 32, 96, 128) -_PREFIXES = (1, 2, 3, 4, 5, 8, 31, 32, 33, 127, 128, 129, 512, 1023, 1024, 1025, 4095, 4096, 8191) +#: requests): both ends of the sort's block kernel, its first cluster-kernel T and +#: the largest capture. The eager prefix lengths straddle the captured ceiling and +#: reach the 8192-row one. +_CAPTURED_ROWS = (1, 4, 5, 128) +_PREFIXES = (1, 4, 5, 127, 128, 129, 4096, 8191) # ── the gate ────────────────────────────────────────────────────────────── @@ -189,10 +190,7 @@ def _weights(gsf: str = "calibrated", layout: str = "cutedsl"): if layout == "shuffle32": perm = _shuffle32(H) w, s = w[:, perm], s[:, perm] - if layout == "linear_weight_scales": - s_k = s.contiguous() - else: - s_k = torch.ops.trtllm.block_scale_interleave(s.contiguous()).view(s.shape) + s_k = torch.ops.trtllm.block_scale_interleave(s.contiguous()).view(s.shape) alpha = (ws2 / GSF[gsf]).contiguous() return w.contiguous().view(torch.float4_e2m1fn_x2), s_k, alpha @@ -369,7 +367,7 @@ def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves th # ── invariance: a pair's row depends on the pair alone ──────────────────── -@pytest.mark.parametrize("tokens", (1, 128, 8192)) +@pytest.mark.parametrize("tokens", (128, 8192)) def test_tiles_and_tactics_agree_per_pair(tokens, certifying): # noqa: F811 """Tile 128 / 256, MMA N 128 / 256, cluster N 1 / 2 and the tuner's fallback all produce every pair's row bit for bit the same, so a cell certifies the @@ -393,13 +391,12 @@ def test_tiles_and_tactics_agree_per_pair(tokens, certifying): # noqa: F811 _assert_pairs_equal(got, first, f"T={tokens} {key} vs (128, 128, 1)") -@pytest.mark.parametrize("tile", (128, 256)) -def test_row_order_and_repeat_do_not_change_a_pair(tile, certifying): # noqa: F811 +def test_row_order_and_repeat_do_not_change_a_pair(certifying): # noqa: F811 """moe_sort orders an expert's rows by atomics (from T = 5). Moving every pair to another row of its expert's run, and calling twice, leaves each pair's row bitwise unchanged: the FC2 result is reproducible run to run.""" certifying(op) - tokens = 4096 + tokens, tile = 4096, 128 ids = _ids("router", tokens, tile, seed=78) sort = _sort(ids, tile) y = _pair_rows(tokens * K, seed=77) @@ -430,12 +427,11 @@ def test_row_order_and_repeat_do_not_change_a_pair(tile, certifying): # noqa: F _assert_pairs_equal(_pairs(out_s, shuffled, tokens), base, "pairs moved within their runs") -@pytest.mark.parametrize("tile", (128, 256)) -def test_eager_rows_prefix_bitwise(tile, certifying): # noqa: F811 +def test_eager_rows_prefix_bitwise(certifying): # noqa: F811 """The 8192-row call is gated against the reference, and every prefix call's pairs are bitwise the 8192-row call's.""" certifying(op) - tokens = 8192 + tokens, tile = 8192, 128 ids = _ids("router", tokens, tile, seed=9) y = _pair_rows(tokens * K, seed=8) sort = _sort(ids, tile) @@ -500,12 +496,11 @@ def _fc1(xq, xsf, sort, tile): ) -@pytest.mark.parametrize("tile", (128, 256)) -def test_fc1_output_feeds_it_at_prefill(tile, certifying): # noqa: F811 +def test_fc1_output_feeds_it_at_prefill(certifying): # noqa: F811 """The real FC1 output at 8192 tokens -- the target's prefill -- passes `_compare`, and the padding rows FC1 zeroes come out zero.""" certifying(op) - tokens = 8192 + tokens, tile = 8192, 128 xq, xsf = _activation(tokens, seed=91) sort = _sort(_ids("router", tokens, tile, seed=92), tile) h, h_sf = _fc1(xq, xsf, sort, tile) @@ -532,13 +527,13 @@ def _capture(fn): return graph, static -@pytest.mark.parametrize("tile", (128, 256)) -def test_captured_rows_replay_on_new_data(tile, certifying): # noqa: F811 +def test_captured_rows_replay_on_new_data(certifying): # noqa: F811 """moe_sort, FC1 and this op are captured together, as the target's graph holds them; each graph is replayed on two fresh activations and id sets, and its written rows are gated against the reference of the graph's own FC1 output and - sort, and bitwise against an eager call on them.""" + sort, and bitwise against an eager call on them. At the target's tile, 128.""" certifying(op) + tile = 128 for tokens in _CAPTURED_ROWS: xq, xsf = _activation(tokens, seed=tokens) ids = _ids("router", tokens, tile, seed=500 + tokens) @@ -579,14 +574,14 @@ def step(xq=xq, xsf=xsf, ids=ids, scales=scales, state=state, tile=tile): # ── what the op writes ──────────────────────────────────────────────────── -@pytest.mark.parametrize("tile", (128, 256)) @pytest.mark.parametrize("tokens", (3, 300)) -def test_unwritten_regions(tokens, tile): +def test_unwritten_regions(tokens): """A sentinel planted in a captured graph's output: which rows replay writes. Every row of the live tiles is written -- the padding rows computed from whatever the input holds there, zero for FC1's zeroed padding -- and the rows past `num_non_exiting_tiles` are not: under a graph they keep whatever the - buffer held.""" + buffer held. Tile 128.""" + tile = 128 args, sort = _problem(tokens, tile, seed=11) graph, out = _capture(lambda: op(**args)) m = out.shape[0] @@ -625,30 +620,21 @@ def test_unwritten_regions(tokens, tile): @pytest.mark.parametrize( "layout", - ( - "shuffle32", - "linear_weight_scales", - "linear_input_scales", - "alpha_rolled", - "alpha_per_layer", - ), + ("shuffle32", "linear_input_scales", "alpha_per_layer"), ) def test_wrong_layouts_are_rejected(layout): - """Each operand prepared the wrong way -- trtllm-gen's 32-row shuffle, weight - scales left linear, the input's scales linear, alpha off by one expert, one - alpha for the whole layer -- returns finite, plausibly scaled output that - `_compare` must reject against the reference of the correctly prepared - operands.""" + """Each operand prepared the wrong way -- the weight in trtllm-gen's 32-row + shuffle, the input's scales linear, one alpha for the whole layer -- returns + finite, plausibly scaled output that `_compare` must reject against the + reference of the correctly prepared operands.""" tokens, tile = 128, 128 good, sort = _problem(tokens, tile, seed=3) m = good["input"].shape[0] - if layout in ("shuffle32", "linear_weight_scales"): + if layout == "shuffle32": bad = _args(good["input"], good["input_scale"], sort, tile, layout=layout) elif layout == "linear_input_scales": linear = unswizzle(good["input_scale"], m, NB).reshape(-1).contiguous() bad = dict(good, input_scale=linear) - elif layout == "alpha_rolled": - bad = dict(good, alpha=good["alpha"].roll(1).contiguous()) else: bad = dict(good, alpha=torch.full_like(good["alpha"], float(good["alpha"].mean()))) out = op(**bad) @@ -679,13 +665,12 @@ def _strided(t: torch.Tensor) -> torch.Tensor: return view.view(t.dtype) if t.element_size() == 1 else view -@pytest.mark.parametrize( - "name", - ("input", "input_scale", "weight", "weight_scale", "alpha", "tile_idx_to_group_idx"), -) +@pytest.mark.parametrize("name", ("input", "weight_scale", "tile_idx_to_group_idx")) def test_noncontiguous_operand_answers_wrongly_and_is_refused(name): """The runner hands data_ptr() to the kernel and never reads a stride: the - same values through a strided view come back finite and wrong.""" + same values through a strided view come back finite and wrong. One activation, + one weight and one sort operand; the guard checks input_scale, weight and + alpha the same way.""" tokens, tile = 64, 128 args, _ = _problem(tokens, tile, seed=21) ref = op.reference(**args) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_weight_path.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_weight_path.py index cd6ffd05d6c1..08e0c2758cd0 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_weight_path.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_weight_path.py @@ -15,9 +15,9 @@ on its own expert, so each expert's output is checked alone: a layout error confined to one expert (a permutation, a scale swizzle, an alpha) cannot hide in a sum over ten. The second is a weighted top-10. Both backends are held to a float64 reference over the dequantized -checkpoint tensors, with `down_proj.input_scale` both spread across the experts and uniform. -When it is uniform the two backends requantize the SwiGLU output identically, so they are -also held to each other. +checkpoint tensors, with `down_proj.input_scale` spread across the experts as a calibrated +checkpoint's is. With a uniform one the two backends requantize the SwiGLU output +identically, so on the weighted routing they are also held to each other. The reference does not model the kernels' accumulation. It models the one rounding step that is not exact: the SwiGLU output is requantized to NVFP4 before FC2, at the backend's @@ -354,16 +354,15 @@ def test_routings_hit_every_expert() -> None: assert torch.unique(ids[weights != 0]).numel() == E, case -@pytest.mark.parametrize("spread", (True, False), ids=["isc2_spread", "isc2_uniform"]) @pytest.mark.parametrize("case", CASES) @pytest.mark.parametrize("backend", BACKENDS) -def test_routed_output_matches_the_float64_reference(backend: str, case: str, spread: bool) -> None: - err = _token_error(_routed(backend, case, spread), _reference(backend, case, spread)) - print(f"[err] {backend} {case} spread={spread}: mean {err.mean():.4e} max {err.max():.4e}") +def test_routed_output_matches_the_float64_reference(backend: str, case: str) -> None: + err = _token_error(_routed(backend, case, True), _reference(backend, case, True)) + print(f"[err] {backend} {case}: mean {err.mean():.4e} max {err.max():.4e}") assert err.max() <= REFERENCE_TOLERANCE, (err.mean().item(), err.max().item()) -@pytest.mark.parametrize("case", CASES) +@pytest.mark.parametrize("case", ("weighted",)) def test_cutedsl_matches_trtllm_gen_under_a_uniform_input_scale(case: str) -> None: err = _token_error(_routed("CUTEDSL", case, False), _routed("TRTLLM", case, False)) print(f"[err] CUTEDSL vs TRTLLM {case}: mean {err.mean():.4e} max {err.max():.4e}") diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py index 971ebb067b05..2fdc11d84ad4 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py @@ -136,7 +136,7 @@ def test_coarse_cell_exercises_tie_breaking(): def test_strided_logits_refused(): # More than one row: a one-row column slice is contiguous by definition. - spec = next(c.spec for c in op.CELLS if c.spec["tokens"] == 8) + spec = next(c.spec for c in op.CELLS if c.spec["tokens"] == 128) args, _, data, sf, _, logits, _, _, _, scalars = _operands(spec, 11) wide = torch.zeros(spec["tokens"], spec["num_experts"] + 8, device=DEV, dtype=logits.dtype) wide[:, : spec["num_experts"]] = logits diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py index 040541c9c4e1..f1df7ac292a4 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py @@ -113,12 +113,12 @@ def test_certified_cells(cell, certifying): # noqa: F811 H._compare(y, exp.to(torch.bfloat16)) -@pytest.mark.parametrize("tokens", [1024, 4096]) -def test_close_to_the_trtllm_gen_entry(tokens): +def test_close_to_the_trtllm_gen_entry(): """Same experts, same routing: within the routed entry's own band of each - other's output (the two differ only in accumulation order).""" - spec = next(c.spec for c in op.CELLS if c.spec["tokens"] == 1024) - spec = {**spec, "tokens": tokens} + other's output (the two differ only in accumulation order). At the row count + where the target switches from one to the other.""" + tokens = 4096 + spec = {**op.CELLS[0].spec, "tokens": tokens} args, ref, data, sf, x, logits, g2 = _operands(spec, 4000 + tokens) y = _cute_dsl(ref, data, sf, logits, g2, spec) s1, sg, s2 = H._scalars(spec["num_experts"], G1, g2) @@ -143,7 +143,7 @@ def test_close_to_the_trtllm_gen_entry(tokens): def test_guard_refuses_a_per_expert_fc2_scale(): - spec = next(c.spec for c in op.CELLS if c.spec["tokens"] == 1024) + spec = {**op.CELLS[0].spec, "tokens": 1024} args, ref, data, sf, x, logits, g2 = _operands(spec, 9) fc1_w, fc1_s, fc2_w, fc2_s = _cute_dsl_experts(ref) s1, sg, s2 = H._scalars(spec["num_experts"], G1, g2) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp8_block_scale_moe_runner.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp8_block_scale_moe_runner.py index dc92b8dd9c6b..d463c4e26622 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp8_block_scale_moe_runner.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp8_block_scale_moe_runner.py @@ -109,7 +109,6 @@ def _args( 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. @@ -128,8 +127,6 @@ def _args( logits = torch.randn(tokens, num_experts, device=DEV, generator=g) weights, ids = torch.topk(torch.softmax(logits, dim=-1), top_k, dim=-1) ids = ids.to(torch.int32) - if foreign_ids: - ids[0, 0], ids[1, 0], ids[2, 0] = -1, num_experts, 40 return dict( routing_logits=None, routing_bias=None, @@ -236,7 +233,7 @@ def _pow2_requantize(x: torch.Tensor) -> torch.Tensor: return (coded.float() * scale).reshape(shp) -@pytest.mark.parametrize("tokens,width", [(16, 256), (64, 640), (256, 128)]) +@pytest.mark.parametrize("tokens,width", [(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 diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fused_sigmoid_gate_mul_add.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fused_sigmoid_gate_mul_add.py index 5aaa3b5107d6..9eed4e2ec169 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fused_sigmoid_gate_mul_add.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fused_sigmoid_gate_mul_add.py @@ -170,47 +170,26 @@ def ratio(wrong: torch.Tensor) -> float: @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", - ], + ["shared_shape_mismatch", "gate_numel_too_small", "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.""" + """Metadata violations are loud, which is why the entry has no guard: one + case each for an operand that disagrees with `final_hidden_states`, a gate + of the wrong length and a destination whose last stride is not 1.""" 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": + if 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) + op(final, gate, shared) def test_cpu_operands_are_rejected_by_triton() -> None: diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_sort.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_sort.py index 9ba068a29ae8..266a8a23a9d1 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_sort.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_sort.py @@ -177,9 +177,9 @@ def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves th # ── CUDA graphs: captured row counts, on new ids ────────────────────────── #: Row counts of a decode capture (requests x MTP verify width 1..4, up to 32 -#: requests): the block kernel (T <= 4), the first cluster-kernel T, and the widths -#: that follow. -_CAPTURED_ROWS = (1, 2, 3, 4, 5, 8, 32, 96, 128) +#: requests): both ends of the block kernel (T <= 4), the first cluster-kernel T, +#: and the largest capture. +_CAPTURED_ROWS = (1, 4, 5, 128) def _capture(args): @@ -196,12 +196,13 @@ def _capture(args): return graph, static_out -@pytest.mark.parametrize("tile", (128, 256)) -def test_captured_rows_replay_on_new_ids(tile, certifying): # noqa: F811 +def test_captured_rows_replay_on_new_ids(certifying): # noqa: F811 """Each graph is replayed on two fresh id sets and gated against the reference, and against an eager call on the same ids: counts and tile arrays bitwise, the - whole output bitwise where the kernel is deterministic (T <= 4).""" + whole output bitwise where the kernel is deterministic (T <= 4). At the + target's tile, 128.""" certifying(op) + tile = 128 for rows in _CAPTURED_ROWS: args = _args(_ids("router", rows, tile, seed=rows), tile, seed=rows) graph, static_out = _capture(args) @@ -226,12 +227,12 @@ def test_captured_rows_replay_on_new_ids(tile, certifying): # noqa: F811 torch.cuda.synchronize() -@pytest.mark.parametrize("tile", (128, 256)) @pytest.mark.parametrize("rows", (3, 300)) -def test_padding_rows_and_spare_tiles_are_not_written(rows, tile): +def test_padding_rows_and_spare_tiles_are_not_written(rows): """A sentinel planted in a captured graph's outputs survives replay in exactly the regions `_compare` excludes; under a graph those regions keep the previous - replay's indices.""" + replay's indices. Block kernel (3 rows) and cluster kernel (300), tile 128.""" + tile = 128 args = _args(_ids("router", rows, tile, seed=5), tile, seed=5) graph, static_out = _capture(args) for trial in (1, 2): @@ -268,15 +269,12 @@ def test_padding_rows_and_spare_tiles_are_not_written(rows, tile): #: Both sides of every kernel-path edge: block / cluster at 4 / 5, the cluster #: kernel's per-block work edges at 102 / 103 and 204 / 205 rows, its fast / slow #: path edges at 3276 / 3277 and 6553 / 6554, and the 8192-row ceiling. -_SWEEP = ( - *(1, 2, 3, 4, 5, 6, 8, 31, 32, 33, 102, 103, 204, 205, 512, 1024), - *(3276, 3277, 4096, 6553, 6554, 8191, 8192), -) +_SWEEP = (4, 5, 102, 103, 204, 205, 3276, 3277, 6553, 6554, 8192) -@pytest.mark.parametrize("tile", (128, 256)) -def test_eager_row_sweep(tile, certifying): # noqa: F811 +def test_eager_row_sweep(certifying): # noqa: F811 certifying(op) + tile = 128 for tokens in _SWEEP: _check(_args(_ids("router", tokens, tile, seed=tokens), tile, seed=tokens)) @@ -287,8 +285,8 @@ def test_eager_row_sweep(tile, certifying): # noqa: F811 def test_block_kernel_rows_follow_token_order(): """T <= 4: each expert's rows in token order, bitwise the reference's.""" for tile in (128, 256): - for tokens in (1, 2, 3, 4): - for seed in range(4): + for tokens in (1, 4): + for seed in range(2): args = _args(_ids("router", tokens, tile, seed=seed), tile, seed=seed) out, ref = _check(args) real = ref[3] >= 0 diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_unpermute_inplace.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_unpermute_inplace.py index d46140fabf49..ef20128ffb78 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_unpermute_inplace.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_unpermute_inplace.py @@ -29,14 +29,12 @@ NAN = float("nan") #: Row counts of a decode capture (requests x MTP verify width 1..4, up to 32 -#: requests), and the eager prefix lengths: both sides of the power-of-two edges up -#: to the 8192-row ceiling. A subset of the prefixes is also re-sorted. -_CAPTURED_ROWS = (1, 2, 3, 4, 5, 8, 32, 96, 128) -_PREFIXES = ( - *(1, 2, 3, 4, 5, 8, 31, 32, 33, 96, 127, 128, 129, 511, 512, 513, 1023, 1024, 1025), - *(2047, 2048, 2049, 4095, 4096, 4097, 8191), -) -_RESORT = (1, 2, 5, 32, 129, 1024, 8192) +#: requests): both ends of the sort's block kernel, its first cluster-kernel T and +#: the largest capture. The eager prefix lengths straddle the captured ceiling and +#: reach the 8192-row one; a subset of them is also re-sorted. +_CAPTURED_ROWS = (1, 4, 5, 128) +_PREFIXES = (1, 4, 5, 127, 128, 129, 4096, 8191) +_RESORT = (1, 5, 129, 8192) def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: @@ -210,16 +208,16 @@ def _capture(fn): return graph, static -@pytest.mark.parametrize("tile", (128, 256)) -def test_captured_rows_replay_on_new_data(tile, certifying): # noqa: F811 +def test_captured_rows_replay_on_new_data(certifying): # noqa: F811 """The sort, the FC2 rows' placement and this op are captured together, as the target runs them; each graph is replayed on two fresh sets of ids, scales and rows. The permuted buffer is not cleared between replays, so the rows the new sort does not use hold the previous placement's rows. Each replay is gated against the reference, and is bitwise equal to an eager call on the same data - through its own sort. + through its own sort. At the target's tile, 128. """ certifying(op) + tile = 128 stale_rows = 0 for tokens in _CAPTURED_ROWS: ids = _ids("router", tokens, seed=tokens) @@ -275,8 +273,7 @@ def step(ids=ids, scales=scales, rows=rows, perm=perm, out=out, tile=tile): # ── eager: prefixes of the largest batch ────────────────────────────────── -@pytest.mark.parametrize("tile", (128, 256)) -def test_eager_rows_prefix_bitwise(tile, certifying): # noqa: F811 +def test_eager_rows_prefix_bitwise(certifying): # noqa: F811 """One block per token makes the grid the only thing T sizes: the 8192-row call is gated against the reference, and the first T rows of every row count in the sweep are bitwise those of the 8192-row call -- over the same placement, and @@ -284,6 +281,7 @@ def test_eager_rows_prefix_bitwise(tile, certifying): # noqa: F811 rows elsewhere. A token's output does not depend on its batch. """ certifying(op) + tile = 128 args, rows, ids = _problem(_spec(8192, tile_tokens_dim=tile), seed=8192 + tile) full = op(**args) _compare(full, op.reference(**args)) @@ -314,7 +312,7 @@ def test_eager_rows_prefix_bitwise(tile, certifying): # noqa: F811 # ── run to run, and where the rows sit ──────────────────────────────────── -@pytest.mark.parametrize("tokens", (1, 32, 128, 8192)) +@pytest.mark.parametrize("tokens", (1, 8192)) def test_repeats_and_placements_are_bitwise(tokens, certifying): # noqa: F811 """What the fused finalize does not give: 8 identical calls, a second sort of the same ids (the sort orders an expert's rows by atomics), a random @@ -478,8 +476,6 @@ def test_op_rejects_loudly(): # ── the guards ──────────────────────────────────────────────────────────── -_OPERANDS = ("permuted_input", "output", "expanded_idx_to_permuted_idx", "topk_scales") - def _strided(t: torch.Tensor, pad: int) -> torch.Tensor: """The same values in a view whose rows are `pad` elements wider apart.""" @@ -488,8 +484,10 @@ def _strided(t: torch.Tensor, pad: int) -> torch.Tensor: return wide[:, : t.shape[1]] -@pytest.mark.parametrize("name", _OPERANDS) +@pytest.mark.parametrize("name", ("permuted_input", "output")) def test_noncontiguous_operand_answers_wrongly_and_is_refused(name): + """One operand the op reads and the one it writes; the guard checks all four + the same way.""" args, _, _ = _problem(_spec(300), seed=61) ref = op.reference(**args) bad = dict(args, **{name: _strided(args[name], 64 if args[name].shape[1] == H else 6)}) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_flashinfer_gemma_rmsnorm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_flashinfer_gemma_rmsnorm.py index a6491f715bd3..0f2d08f82828 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_flashinfer_gemma_rmsnorm.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_flashinfer_gemma_rmsnorm.py @@ -34,14 +34,9 @@ def randn(*s): rows, off = shape[0], spec["col_offset"] buffer = randn(rows, spec["row_width"]) x = buffer[:, off : off + width] - elif "row_step" in spec: - buffer = randn(shape[0] * spec["row_step"], width) - x = buffer[:: spec["row_step"]] else: buffer = x = randn(*shape) - weight = (torch.randn(width, generator=g, device="cuda") * spec.get("weight_scale", 1.0)).to( - dtype - ) + weight = torch.randn(width, generator=g, device="cuda").to(dtype) return buffer, x, weight, spec["eps"] diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_layer_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_layer_norm.py index e297e7674808..921abba692ca 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_layer_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_layer_norm.py @@ -50,31 +50,21 @@ def _affine(gen, width=WIDTH, dtype=torch.bfloat16): def _build(spec, seed): gen = _gen(seed) shape, dtype = spec["shape"], spec["dtype"] - if spec.get("layout", "contiguous") == "column_slice": - rows, width = shape - wide = torch.randn(rows, 2 * width, generator=gen, device=DEVICE).to(dtype) - x = wide[:, width:] - assert not x.is_contiguous() and x.stride(-1) == 1 - else: - x = torch.randn(shape, generator=gen, device=DEVICE).to(dtype) - weight, bias = _affine(gen, width=shape[-1], dtype=spec.get("param_dtype", dtype)) - return x, weight, (bias if spec.get("bias", True) else None), EPS + x = torch.randn(shape, generator=gen, device=DEVICE).to(dtype) + weight, bias = _affine(gen, width=shape[-1], dtype=dtype) + return x, weight, bias, EPS @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) def test_certified_cells(cell, certifying): # noqa: F811 certifying(op) x, weight, bias, eps = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) - before = [t.clone() for t in (x, weight, bias) if t is not None] + before = [t.clone() for t in (x, weight, bias)] out = op(x, weight, bias, eps) assert out.shape == x.shape and out.dtype == x.dtype assert out.is_contiguous() and out.data_ptr() != x.data_ptr() - after = [t for t in (x, weight, bias) if t is not None] + after = [x, weight, bias] assert all(torch.equal(a, b) for a, b in zip(before, after)), "inputs must not be mutated" - if x.dim() == 3: - assert torch.equal(out.reshape(-1, WIDTH), op(x.reshape(-1, WIDTH), weight, bias, eps)) - if not x.is_contiguous(): - assert torch.equal(out, op(x.contiguous(), weight, bias, eps)) assert torch.equal(out, op(x, weight, bias, eps)), "two launches must agree bitwise" torch.testing.assert_close(out, op.reference(x, weight, bias, eps)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_rms_norm_gated_token_major.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_rms_norm_gated_token_major.py index 5d1725380046..ca8bbbbb80ac 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_rms_norm_gated_token_major.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_rms_norm_gated_token_major.py @@ -29,10 +29,12 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: 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 #: The target's concatenated GDN input projection: [q k v | z | a | b] = 10240 | 6144 | 48 | 48. -_IN_PROJ_WIDTH = 16480 +_IN_PROJ_WIDTH, _Z_OFFSET = 16480, 10240 + + +#: The checkpoint's GDN epilogue geometry over a packed `z`, for the non-cell tests. +_PACKED = dict(op.CELLS[0].spec, tokens=17, z="packed") def _seed(text: str) -> int: @@ -47,21 +49,13 @@ def _build(spec, seed): def randn(*shape): return torch.randn(*shape, generator=g, device="cuda").to(dtype) - if spec.get("x") == "strided_rows": - x = randn(t * heads, 2 * n)[:, :n] - else: - x = randn(t * heads, n) - if "x_scale" in spec: - x = x * spec["x_scale"] + x = randn(t * heads, n) layout = spec.get("z", "packed") - if layout == "column_slice": - z = randn(t, _PROJ_WIDTH)[:, _Z_OFFSET : _Z_OFFSET + heads * n].view(t, heads, n) - elif layout == "in_proj": + if layout == "in_proj": z = randn(t, _IN_PROJ_WIDTH)[:, _Z_OFFSET : _Z_OFFSET + heads * n].view(t, heads, n) - elif layout == "head_strided": - z = randn(t, heads, 2 * n)[:, :, :n] else: + assert layout == "packed", layout z = randn(t, heads, n) weight = randn(n) @@ -76,14 +70,8 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res 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("z") == "in_proj" and cell.spec["tokens"] > 1: assert z.stride() == (_IN_PROJ_WIDTH, cell.spec["width"], 1) 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) @@ -98,10 +86,9 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res _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: +def test_fp8_byte_gate_discriminates() -> None: """Each wrong variant of the reference misses the byte gate on most elements.""" - spec = dict(op.CELLS[-1].spec, fp8_scale=scale_value) + spec = dict(op.CELLS[-1].spec, fp8_scale=0.0234375) 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)) @@ -112,18 +99,17 @@ def test_fp8_byte_gate_discriminates(scale_value: float) -> None: "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) + 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) + """A kernel ignoring `gate_activation` would pass the sigmoid cells; the two gates differ.""" + x, z, weight, eps, _, _ = _build(_PACKED, seed=3) assert not torch.equal( op(x, z, weight, eps, None, "silu"), op(x, z, weight, eps, None, "sigmoid") ) @@ -131,7 +117,7 @@ def test_the_two_gates_differ() -> None: def test_guard_refuses_a_weight_of_the_wrong_length() -> None: """Unguarded, a length-192 weight is silently truncated to weight[:128].""" - x, z, weight, eps, _, act = _build(dict(op.CELLS[1].spec, tokens=4), seed=12) + x, z, weight, eps, _, act = _build(dict(_PACKED, 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), @@ -147,7 +133,7 @@ def test_guard_refuses_a_weight_of_the_wrong_length() -> None: 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) + x, z, weight, eps, _, act = _build(dict(_PACKED, 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") @@ -166,13 +152,13 @@ def test_guard_refuses_a_malformed_fp8_scale() -> None: def test_rejects_unsupported_gate_activation() -> None: - x, z, weight, eps, _, _ = _build(dict(op.CELLS[1].spec, tokens=4), seed=10) + x, z, weight, eps, _, _ = _build(dict(_PACKED, 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) + x, z, weight, eps, _, act = _build(dict(_PACKED, 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"): @@ -182,7 +168,7 @@ def test_rejects_z_that_does_not_cover_x() -> None: def test_prefill_launch_rejects_z_that_does_not_cover_x() -> None: """From 32768 rows the op runs its own multi-row launch; a `z` one token short still reaches the registered op's check instead of being read past its end.""" - x, z, weight, eps, _, act = _build(dict(op.CELLS[1].spec, tokens=683), seed=12) + x, z, weight, eps, _, act = _build(dict(_PACKED, tokens=683), seed=12) assert x.shape[0] >= 32768, x.shape with pytest.raises(AssertionError, match="does not match x shape"): op(x, z[:-1], weight, eps, None, act) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_gated_ngram_injection.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_gated_ngram_injection.py index e08dc99a1dfe..c410d9e790ab 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_gated_ngram_injection.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_gated_ngram_injection.py @@ -112,12 +112,10 @@ def _build(spec, seed): dots = spec["dots"] key, query, value = _single_channel_operands(dots, spec["value_dtype"]) return key, query, value, len(dots), spec["clamp_min"] - hc_count = spec.get("hc_count", HC_COUNT) - dtype = spec.get("dtype", torch.bfloat16) key, query, value = _make_operands( - seed, dtype, spec["leading"], hc_count, spec.get("value_dtype") + seed, torch.bfloat16, spec["leading"], value_dtype=spec.get("value_dtype") ) - return key, query, value, hc_count, spec.get("clamp_min", CLAMP_MIN) + return key, query, value, HC_COUNT, spec.get("clamp_min", CLAMP_MIN) @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) @@ -144,7 +142,7 @@ def test_flattened_call_matches_batched() -> None: 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): + for scale in (1.0, 4096.0): got = op( (key.float() * scale).to(torch.bfloat16), (query.float() * scale).to(torch.bfloat16), diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ngram_token_history.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ngram_token_history.py index b6069fe13f5b..680aa9320e7c 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ngram_token_history.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ngram_token_history.py @@ -69,39 +69,27 @@ def _qsl(lengths, dtype, qsl_from="literal"): def _build(spec, seed): - """`(args, backing_pool)`: the op's positional arguments, and the tensor - the pool argument views (itself, unless `layout == "pool_view"`).""" + """The op's positional arguments for one cell.""" g = torch.Generator().manual_seed(seed) - ngram_size = spec.get("ngram_size", NGRAM) lengths = spec["lengths"] total = sum(lengths) ids = torch.randint(0, CKPT_VOCAB_SIZE, (total,), generator=g, dtype=torch.long) density = spec.get("eos_density", 0.0) if density > 0.0: ids[torch.rand(total, generator=g) < density] = EOS - layout = spec.get("layout", "contiguous") - pool_rows = spec["pool_slots"] * (2 if layout == "pool_view" else 1) - backing = torch.randint( - 0, CKPT_VOCAB_SIZE, (pool_rows, ngram_size - 1), generator=g, dtype=torch.long + pool = torch.randint( + 0, CKPT_VOCAB_SIZE, (spec["pool_slots"], CONTEXT_LEN), generator=g, dtype=torch.long ).to(DEV) - pool = backing[::2] if layout == "pool_view" else backing - input_ids = ids.to(DEV, spec.get("ids_dtype", torch.int64)) - if layout == "strided_ids": - wide = torch.zeros(2 * total, dtype=input_ids.dtype, device=DEV) - wide[::2] = input_ids - input_ids = wide[::2] - assert input_ids.stride(0) == 2 index_dtype = spec.get("index_dtype", torch.int32) - args = ( - input_ids, + return ( + ids.to(DEV, spec.get("ids_dtype", torch.int64)), pool, torch.tensor(spec["slots"], dtype=index_dtype, device=DEV), torch.tensor(spec["carries"], dtype=torch.bool, device=DEV), _qsl(lengths, spec.get("qsl_dtype", torch.int32), spec.get("qsl_from", "literal")), EOS, - ngram_size, + NGRAM, ) - return args, backing def _expect(args): @@ -112,9 +100,8 @@ def _expect(args): @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name certifying(op) - args, backing = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) read_only = [a.clone() for a in (args[0], args[2], args[3], args[4])] - backing_before = backing.clone() pool_before = args[1].clone() expected = _expect(args) window = op(*args) @@ -124,8 +111,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res # The whole-pool comparison locks the unnamed slots; this keeps it from # being vacuous -- every cell has a sequence that publishes new ids. assert not torch.equal(args[1], pool_before), "no slot moved" - if cell.spec.get("layout") == "pool_view": - assert torch.equal(backing[1::2], backing_before[1::2]), "the view wrote outside itself" # ----------------------------------------------------------- multi-step @@ -236,7 +221,7 @@ def test_wrong_variants_are_separated() -> None: pool_slots=6, eos_density=0.08, ) - args, _ = _build(spec, seed=9001) + args = _build(spec, seed=9001) ids, pool, slots, carries, qsl = args[:5] start = pool.clone() window = op(*args) @@ -342,12 +327,10 @@ def test_guard_refuses_a_short_query_start_loc() -> None: "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"], + ids=["pool_too_wide", "int_flags", "ngram_size_1"], ) def test_loud_refusals_are_not_guarded(index, value, match) -> None: args = list(_two_sequences()) @@ -382,7 +365,7 @@ def test_cuda_graph_capture_and_replay() -> None: ids_dtype=torch.int32, qsl_dtype=torch.int64, ) - args, _ = _build(spec, seed=11001) + args = _build(spec, seed=11001) start = args[1].clone() side = torch.cuda.Stream() diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_gate_conv_inject_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_gate_conv_inject_norm.py index 0cfbdf6de72c..3949c1e76bec 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_gate_conv_inject_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_gate_conv_inject_norm.py @@ -104,9 +104,7 @@ def _build(spec, seed): g = torch.Generator(device=DEV).manual_seed(seed) gc = torch.Generator().manual_seed(seed) lengths = spec["lengths"] - c = spec.get("hc_count", CKPT_HC_COUNT) - d = spec.get("hidden", CKPT_HIDDEN) - dtype = spec.get("dtype", torch.bfloat16) + c, d, dtype = CKPT_HC_COUNT, CKPT_HIDDEN, torch.bfloat16 w = c * d rows = sum(lengths) num_seqs = len(lengths) @@ -119,18 +117,12 @@ def rand(*shape, scale=1.0): scales = torch.tensor([0.5 * 2.0 ** (i % 4) for i in range(c)], device=DEV) hc = (rand(rows, w, scale=3.0).float() * scales.repeat_interleave(d)).to(dtype) key = rand(rows, w) - if spec.get("layout") == "row_slice": - hbuf, kbuf = rand(2 * rows, w), rand(2 * rows, w) - hbuf[::2], kbuf[::2] = hc, key - hc, key = hbuf[::2], kbuf[::2] value = rand(rows, d, scale=0.5) norms = [rand(w, scale=GAMMA_RMS) for _ in range(4)] conv_weight = rand(w, 1, 4, scale=0.3) conv_states = rand(pool_slots, w, 3 * CKPT_DILATION) order = (torch.randperm(pool_slots - 2, generator=gc) + 1).tolist() slots = ([0, pool_slots - 1] + order)[:num_seqs] if pool_slots >= 64 else list(range(num_seqs)) - if "pad" in spec: - slots[spec["pad"]] = PAD slots = torch.tensor(slots, dtype=torch.int32, device=DEV) starts = [0] for length in lengths: @@ -220,14 +212,11 @@ def _old( @pytest.mark.parametrize( "lengths,carries", [ - ((1,), None), - ((1,) * 32, None), ((1,) * 128, None), ((4,) * 8, None), - ((8192,), (False,)), ((5000, 3000) + (1,) * 32, (False, True) + (True,) * 32), ], - ids=["m1", "m32", "m128", "mtp4x8", "prefill8192", "mixed2ctx32dec"], + ids=["m128", "mtp4x8", "mixed2ctx32dec"], ) def test_matches_the_composition_it_replaces(lengths, carries) -> None: args = _build(dict(lengths=lengths, pool_slots=461, carries=carries), seed=len(lengths)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_ngram_embed.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_ngram_embed.py index ea2c43621f6c..678e021561b9 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_ngram_embed.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_ngram_embed.py @@ -62,20 +62,15 @@ def _scattered(n: int, pool_slots: int, g) -> list: def _geometry(spec, g): - """`(multipliers, vocab, offsets)` for the cell: the checkpoint's own for a - 3-gram of 8 heads per order, deterministic stand-ins otherwise.""" - n = spec.get("ngram_size", 3) - hp = spec.get("heads_per_ngram", 8) - if n == 3 and hp == 8 and spec.get("table") == "ckpt": + """`(multipliers, vocab, offsets)` for the cell: the checkpoint's own over + the full-height table, the checkpoint's multipliers over small stand-in head + vocabularies otherwise.""" + if spec.get("table") == "ckpt": return CKPT_MULTIPLIERS, CKPT_HEAD_VOCAB_SIZES, CKPT_HEAD_OFFSETS - mult = ( - CKPT_MULTIPLIERS - if n == 3 - else tuple(int(x) | 1 for x in torch.randint(10**12, 3 * 10**13, (n,), generator=g)) - ) + mult = CKPT_MULTIPLIERS vocab = tuple( 1000 + 7 * h + int(x) - for h, x in enumerate(torch.randint(0, 50, (hp * (n - 1),), generator=g)) + for h, x in enumerate(torch.randint(0, 50, (len(CKPT_HEAD_VOCAB_SIZES),), generator=g)) ) offsets, acc = [], 0 for v in vocab: @@ -107,12 +102,12 @@ def _with_eos(x: torch.Tensor, density: float, g) -> torch.Tensor: def _build(spec, seed): g = torch.Generator().manual_seed(seed) lengths = spec["lengths"] - n = spec.get("ngram_size", 3) + n = len(CKPT_MULTIPLIERS) num_seqs = len(lengths) pool_slots = spec.get("pool_slots", num_seqs + 2) density = spec.get("eos_density", 0.0) ids = _with_eos(torch.randint(0, VOCAB, (sum(lengths),), generator=g), density, g) - ids = ids.to(spec.get("ids_dtype", torch.int32)).to(DEV) + ids = ids.to(torch.int32).to(DEV) pool = _with_eos(torch.randint(0, VOCAB, (pool_slots, n - 1), generator=g), density, g).to(DEV) slots = _scattered(num_seqs, pool_slots, g) if pool_slots >= 64 else list(range(num_seqs)) slots = torch.tensor(slots, dtype=torch.int32, device=DEV) @@ -174,8 +169,6 @@ def _old(ids, pool, slots, carries, qsl, table, scale, mult, vocab, offsets, eos [ dict(lengths=(1,) * 128, pool_slots=ENGINE_SLOTS, eos_density=0.05), dict(lengths=(4,) * 32, pool_slots=ENGINE_SLOTS, eos_density=0.05), - dict(lengths=(1,) * 8, table="host"), - dict(lengths=(8192,), carries=(False,), pool_slots=ENGINE_SLOTS, eos_density=0.01), dict( lengths=(5000, 3000) + (1,) * 32, carries=(False, True) + (True,) * 32, @@ -183,7 +176,7 @@ def _old(ids, pool, slots, carries, qsl, table, scale, mult, vocab, offsets, eos eos_density=0.05, ), ], - ids=["decode128", "mtp4x32", "host8", "prefill8192", "mixed2ctx32dec"], + ids=["decode128", "mtp4x32", "mixed2ctx32dec"], ) def test_matches_the_calls_it_replaces(spec) -> None: args = _build(spec, seed=len(spec["lengths"])) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_dequantize_e4m3_per_tensor.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_dequantize_e4m3_per_tensor.py index a58ee2bdde8b..0609b96981c0 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_dequantize_e4m3_per_tensor.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_dequantize_e4m3_per_tensor.py @@ -81,26 +81,12 @@ def _input(spec, seed: int) -> torch.Tensor: 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]) @@ -114,7 +100,7 @@ def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves th certifying(op) x = _input(cell.spec, zlib.crc32(cell.why.encode())) for value in cell.spec["scales"]: - scales = _scales(cell.spec, value, x.dim()) + scales = _scale(value, x.dim(), cell.spec["dtype"]) out = op(x, scales) assert out.shape == x.shape and out.dtype == scales.dtype and out.is_contiguous() _compare(out, op.reference(x, scales)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_fp8_quantize_1x128.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_fp8_quantize_1x128.py index 2a65ca954ee7..fdad8d2dfddc 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_fp8_quantize_1x128.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_fp8_quantize_1x128.py @@ -144,10 +144,10 @@ def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves th # ── 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): +def test_use_ue8m0_is_bit_exact_and_a_different_quantization(): """Once the dequant scale is a power of two its reciprocal is exact, so the approximate division cannot reach the result.""" + tokens, hidden = 17, 2560 x = _input(tokens, hidden, seed=5 + tokens) data, sf = op(x, True) ref_data, ref_sf = op.reference(x, True) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_fwd.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_fwd.py index 2e691b0f38fe..e4bcbccce9e2 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_fwd.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_fwd.py @@ -169,56 +169,32 @@ def test_the_initial_state_is_read() -> None: assert torch.equal(with_carry[:, 3:], without[:, 3:]) +#: One refusal per family the op checks itself: a kernel/layout mismatch, a +#: metadata dtype, an operand dtype. _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.""" + """What 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": + if 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) @@ -228,10 +204,8 @@ 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 - ) + if case == "width_five": + spec.update(width=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": @@ -246,9 +220,7 @@ def _guard_case(case: str): return args -@pytest.mark.parametrize( - "case", ["width_one", "width_five", "weight_dtype", "token_stride", "state_stride"] -) +@pytest.mark.parametrize("case", ["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.""" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_triton_update.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_triton_update.py index 110c24a273ed..53881dbb4a3c 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_triton_update.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_triton_update.py @@ -205,25 +205,20 @@ def test_a_short_pool_writes_into_the_next_slot_and_is_refused() -> None: @pytest.mark.parametrize( - "case,overrides,match", + "case,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"), + # The pool-width guard; the weight shares it. + ("narrow_pool", "must cover all 512 channels"), + ("short_snapshot_buffer", "intermediate_conv_window must be at least"), ], ) -def test_guard_refuses_undersized_buffers(case, overrides, match) -> None: +def test_guard_refuses_undersized_buffers(case, 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): diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_update.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_update.py index 9f468a8640b9..baa1f27bcfaa 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_update.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_update.py @@ -138,63 +138,40 @@ def _small(**overrides): return _build(spec, seed=14)[0] +#: One metadata violation per family the op checks: the conv width, the slot +#: indices, the operand dtypes. _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.""" +def test_op_rejects_metadata_violations_loudly(case: str) -> None: + """Why the entry has no `is_valid`: the op raises on metadata violations + itself (one probe per family here).""" fragment, overrides = _REJECTIONS[case] args = _small(**(overrides or {})) - if case == "indices_length": - args[6] = args[6][:3] - elif case == "indices_int64": + if 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) -@pytest.mark.parametrize("rows", [64, 121, 128]) -def test_decode_rows_match_the_trtllm_kernel_bitwise(rows: int) -> None: - """From 64 rows the decode shape runs the CuTe kernel: x and the whole pool - bitwise what `torch.ops.trtllm.causal_conv1d_update` leaves, a pad row untouched. +def test_decode_rows_match_the_trtllm_kernel_bitwise() -> None: + """From 64 rows the decode shape runs the CuTe kernel: at that boundary x and + the whole pool are bitwise what `torch.ops.trtllm.causal_conv1d_update` leaves, + a pad row untouched. x is the q|k|v column view of an in_proj output; the pool a per-layer view of a two-layer buffer, as the engine hands both over.""" + rows = 64 g = torch.Generator(device="cuda").manual_seed(rows) dim = 10240 proj = torch.randn(rows, dim + _IN_PROJ_TAIL, generator=g, device="cuda").to(torch.bfloat16) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_chunk_gated_delta_rule.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_chunk_gated_delta_rule.py index 678033a7a15a..0c389d3a15b0 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_chunk_gated_delta_rule.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_chunk_gated_delta_rule.py @@ -5,7 +5,7 @@ The first indexed in-place call per autotune key in a process must be as correct as a warm one: the state kernel's autotuner restores the indexed state its benchmark launches write, and the pool's dtype is part of the key. "First call in the process" is a property -of the process, so the cold-call tests re-run this file as a child: +of the process, so the cold-call test re-runs this file as a child: `python `. """ @@ -79,8 +79,6 @@ def _operands( 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, @@ -94,29 +92,27 @@ def _operands( 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) + 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 + 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) + 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() + 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, dtype: torch.dtype = torch.float32 -) -> torch.Tensor: +def _pool(num_slots: int, seed: int, dtype: torch.dtype = torch.float32) -> torch.Tensor: """A recurrent-state pool, `[slots, H, V, K]` (K innermost), at `dtype`.""" gen = torch.Generator(device=DEV).manual_seed(seed) - shape = (num_slots, num_v_heads, HEAD_V_DIM, HEAD_K_DIM) + shape = (num_slots, NUM_V_HEADS, HEAD_V_DIM, HEAD_K_DIM) return (torch.randn(shape, generator=gen, device=DEV) * 0.1).to(dtype) @@ -288,8 +284,8 @@ def test_gate_discriminates_wrong_variants() -> None: # --------------------------------------------------------------------------- -# The cold call. The cells above have already benchmarked these autotune keys by -# the time pytest reaches these, so each runs in a fresh interpreter. +# The cold call. The cells above have already benchmarked this autotune key by +# the time pytest reaches it, so it runs in a fresh interpreter. CHILD_TIMEOUT_S = 420 #: The GDN prefill call shape, packed varlen, state by slot index. @@ -331,25 +327,8 @@ def _child_cold_call(pool_dtype: torch.dtype = torch.float32) -> None: _gate_pool(args[6], ref) -def _child_cold_call_odd_heads() -> None: - """A head geometry (5 key / 20 value heads) nothing else uses.""" - autotuner = _state_kernel_autotuner() - assert not autotuner.cache, "the child must be the first indexed call in its process" - q, k, v, g, beta, cu = _operands((96,), seed=83, num_k_heads=5, num_v_heads=20) - pool = _pool(4, seed=831, num_v_heads=20) - indices = torch.tensor([1], dtype=torch.int32, device=DEV) - args = (q, k, v, g, beta, None, pool, indices, True, False, cu, False, None) - ref = op.reference(*args) - o, _ = op(*args) - assert len(autotuner.cache) == 1, "the call must have benchmarked exactly one key" - _gate_output(o, ref) - _gate_pool(pool, ref) - - _CHILD_MODES = { - "cold-call": _child_cold_call, "cold-call-bf16-pool": lambda: _child_cold_call(torch.bfloat16), - "cold-call-odd-heads": _child_cold_call_odd_heads, } @@ -373,21 +352,12 @@ def _run_child(mode: str) -> None: ) -def test_cold_indexed_call_at_this_checkpoints_geometry() -> None: - """The first indexed call of a process at 16 / 48 heads, K = V = 128, in the - mixer's shape (256 + 96 packed, by slot, into `output=`), matches the reference.""" - _run_child("cold-call") - - def test_cold_indexed_call_over_a_bf16_pool() -> None: - """The same call over a bf16 pool, the autotune key an fp32 pool never touches.""" + """The first indexed call of a process at 16 / 48 heads, K = V = 128, over a bf16 + pool, in the mixer's shape (256 + 96 packed, by slot, into `output=`), matches the + reference.""" _run_child("cold-call-bf16-pool") -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 index 5fcd1d61448b..27a365eaf86c 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_cu_seqlens_from_seq_lens.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_cu_seqlens_from_seq_lens.py @@ -49,12 +49,6 @@ def _build(spec): buf = torch.full((len(lengths) + 9,), 99, dtype=dtype, device=DEV) buf[: len(lengths)] = _seq_lens(lengths, dtype) return buf[: len(lengths)] - if layout == "strided": - buf = torch.full((2 * len(lengths),), 99, dtype=dtype, device=DEV) - buf[::2] = _seq_lens(lengths, dtype) - view = buf[::2] - assert not view.is_contiguous() - return view return _seq_lens(lengths, dtype) @@ -132,11 +126,11 @@ def test_accumulation_happens_in_the_requested_dtype() -> None: assert op(seq_lens, torch.int32).tolist() == [0, 2_000_000_000, -294_967_296] -@pytest.mark.parametrize("batch, p", ((8, 2), (32, 8))) -def test_cuda_graph_capture_and_replay(batch, p) -> None: +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. Eight - requests of 2 rows, and 32 of 8 (the MTP-7 verify's 256 rows).""" + reaches Python, the values are read on device and never synchronized on. + 32 requests of 8 rows (the MTP-7 verify's 256 rows).""" + batch, p = 32, 8 static_lens = _seq_lens([p] * batch) side = torch.cuda.Stream() side.wait_stream(torch.cuda.current_stream()) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_dilated_short_conv_paged.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_dilated_short_conv_paged.py index 7233481e3ef2..6dbc9a0eafdc 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_dilated_short_conv_paged.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_dilated_short_conv_paged.py @@ -333,10 +333,10 @@ def test_repeat_calls_are_bitwise_stable() -> None: assert torch.equal(first[0], again[0]) and torch.equal(first[1], again[1]) -@pytest.mark.parametrize("batch, p", ((4, 4), (32, 8))) -def test_cuda_graph_capture_and_replay(batch, p) -> None: - """Captured once, replayed after the batch tensors' values change in place: four - requests of 4 tokens, and the MTP-7 verify's 32 of 8.""" +def test_cuda_graph_capture_and_replay() -> None: + """Captured once at the MTP-7 verify's 32 requests of 8 tokens, replayed after + the batch tensors' values change in place.""" + batch, p = 32, 8 args = _small((p,) * batch, tuple(range(batch)), seed=25, pool_slots=batch + 2) x, weight, _, pool, qsl, cidx, his = args[:7] live = pool.clone() @@ -372,10 +372,9 @@ def test_cuda_graph_capture_and_replay(batch, p) -> None: @pytest.mark.parametrize( "case,fragment", [ + # One per torch call that refuses: the pool's index copy, the convolution. ("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: @@ -386,10 +385,6 @@ def test_mismatched_operands_raise(case: str, fragment: str) -> None: 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) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_extract_transpose_prefill_slice.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_extract_transpose_prefill_slice.py index 233386dcf207..cf25a02c5dcf 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_extract_transpose_prefill_slice.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_extract_transpose_prefill_slice.py @@ -30,12 +30,7 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: def _src(num_rows: int, num_cols: int, dtype: torch.dtype, seed: int) -> torch.Tensor: g = torch.Generator(device="cuda").manual_seed(seed) - base = torch.randn(num_rows, num_cols, generator=g, device="cuda") - if dtype in (torch.int32, torch.int8): - return (base * 10).to(dtype) - if dtype is torch.bool: - return base > 0 - return base.to(dtype) + return torch.randn(num_rows, num_cols, generator=g, device="cuda").to(dtype) def _build(spec, seed): @@ -74,9 +69,8 @@ def test_the_callable_refuses_a_non_unit_column_stride() -> None: [ (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"], + ids=["rows_past_src", "cols_past_src"], ) 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 diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_gating.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_gating.py index 38a6cf3f54d1..1a77c6d35c26 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_gating.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_gating.py @@ -29,7 +29,8 @@ assert torch.cuda.is_available(), "fused_gdn_gating requires a CUDA device" -#: The smallest softplus `beta` a cell drives; the tail error scales as 1 / beta. +#: The softplus `beta` the absolute floor is sized for; the tail error scales as +#: 1 / beta. The target calls 1.0, so 0.5 leaves a factor 2 of headroom. _MIN_CERTIFIED_BETA = 0.5 @@ -38,9 +39,9 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: relative term is `4 * 2^-22`. The absolute floor is the unstabilised softplus's tail error -- about one approximate ulp of the log's argument -- scaled by `1 / beta` and by the `exp(A_log)` it multiplies; - `_compare` cannot see either, so it takes the certified extremes - (exp(A_log) <= 139.42, beta >= 0.5): atol = 4 * 2^-22 * 139.42 / 0.5 = - 2.66e-04. + `_compare` cannot see either, so it takes exp(A_log) <= 139.42 (the + checkpoint's extreme) and beta >= 0.5 (`_MIN_CERTIFIED_BETA`): atol = + 4 * 2^-22 * 139.42 / 0.5 = 2.66e-04. Measured against this gate on 256 checkpoint-range rows: worst error / allowance 0.215. Five wrong formulas (dt_bias dropped, sign dropped, @@ -71,7 +72,7 @@ def _params(g, num_heads=NUM_V_HEADS, dtype=torch.float32): def _build(spec, seed): """`__call__`'s positional arguments for one spec.""" g = torch.Generator(device=DEV).manual_seed(seed) - rows, heads = spec["rows"], spec.get("num_heads", NUM_V_HEADS) + rows, heads = spec["rows"], NUM_V_HEADS a_dtype = spec.get("a_dtype", torch.bfloat16) layout = spec.get("layout", "contiguous") @@ -84,18 +85,9 @@ def randn(*shape): # `_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]) + A_log, dt_bias = _params(g, heads) return A_log, a, dt_bias, spec.get("beta", 1.0), spec.get("threshold", SOFTPLUS_THRESHOLD) @@ -106,14 +98,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res g = op(*args) _compare(g, op.reference(*args)) assert g.is_contiguous() and torch.isfinite(g).all() - A_log, a, dt_bias = args[:3] - if not a.is_contiguous(): - # The row stride is consumed in place, not approximated. - assert torch.equal(g, op(A_log, a.contiguous(), dt_bias, *args[3:])) - if "overlong" in cell.spec: - heads = a.shape[1] - trimmed = op(A_log[:heads].clone(), a, dt_bias[:heads].clone(), *args[3:]) - assert torch.equal(g, trimmed), "elements past a.shape[1] must not be read" def test_threshold_branch_is_live_for_this_checkpoint() -> None: @@ -183,19 +167,12 @@ def test_gate_discriminates_wrong_variants() -> None: 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: +def test_op_rejects_a_non_2d_a() -> 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): + a = torch.randn(2, 4, NUM_V_HEADS, device=DEV).to(torch.bfloat16) + with pytest.raises(ValueError, match="too many values to unpack"): with validating(op): op(A_log, a, dt_bias) @@ -220,12 +197,8 @@ def _guard_case(case): 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]) + return (wide[:40], a, dt_bias), (wide[:NUM_V_HEADS], a, dt_bias) @pytest.mark.parametrize( @@ -233,9 +206,7 @@ def _guard_case(case): [ ("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: diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_post_conv.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_post_conv.py index 325bb8925910..d4280185fe90 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_post_conv.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_post_conv.py @@ -75,9 +75,8 @@ def _compare(out: tuple, ref: tuple) -> None: DTYPE = torch.bfloat16 HEADS = (NUM_K_HEADS, HEAD_K_DIM, NUM_V_HEADS, HEAD_V_DIM) L2_NORM_EPS = 1e-6 -#: The stock mixer's `z` half of `in_proj_qkvz`, and its `[b|a]` projection width. +#: The `z` columns of the input projection. Z_WIDTH = NUM_V_HEADS * HEAD_V_DIM -BA_WIDTH = 2 * NUM_V_HEADS #: The target's one GDN input-projection output: [q|k|v | z | a | b]. IN_PROJ_WIDTH = QKV_DIM + Z_WIDTH + 2 * NUM_V_HEADS # 16480 @@ -92,13 +91,7 @@ def _build(spec, seed): 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 == "in_proj": + if layout == "in_proj": # `_gdn_block`: one [T, 16480] GEMM output; q|k|v, a and b are its column # views, the decode half the row split of q|k|v, the prefill half the # transposing extraction of the rest. @@ -117,11 +110,6 @@ def randn(*shape): 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 diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py index b47fd95fa910..341f11ff8c21 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py @@ -525,9 +525,9 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res # ── CUDA graphs ─────────────────────────────────────────────────────────── -_GRAPH_CASES = [(n, t, torch.float32) for t in range(2, 9) for n in (1, 8, 15, 16, 32)] + [ - (n, t, torch.bfloat16) for t in (2, 4, 8) for n in (1, 15, 16, 32) -] +#: The in-kernel commit at one sequence and the smallest token tile, the all-layer commit at +#: the envelope's 32 requests and the largest, and a bf16 checkpoint at the partition threshold. +_GRAPH_CASES = [(1, 2, torch.float32), (32, 8, torch.float32), (16, 4, torch.bfloat16)] @pytest.mark.parametrize( @@ -536,11 +536,9 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res ids=[f"{n}-{t}-{'fp32' if d is torch.float32 else 'bf16'}" for n, t, d in _GRAPH_CASES], ) def test_cuda_graph_replay_on_new_data_is_bitwise_eager(n, t, state_dtype) -> None: - """Every T 2..8 (MTP draft lengths 1..7) over an fp32 checkpoint, and T 2, 4 and 8 over a - bf16 one, at the batches either side of the partition threshold and at 1, 8 and 32: two - layers captured with, from 16 sequences on, the all-layer commit after them; replayed - three times on new activations, with acceptance in between, against the same steps run - eagerly on a twin state.""" + """Two layers captured with, from 16 sequences on, the all-layer commit after them; + replayed three times on new activations, with acceptance in between, against the same + steps run eagerly on a twin state.""" assert ALL_LAYER_COMMIT_MIN_BATCH == CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE path = "all_layer" if n >= ALL_LAYER_COMMIT_MIN_BATCH else "small" spec = dict(path=path, n=n, t=t, layers=2, slots=40, state_dtype=state_dtype) @@ -614,25 +612,24 @@ def _drive(bench, fault): return None -@pytest.mark.parametrize("t", (4, 8)) @pytest.mark.parametrize( "fault", ("none", "commit skipped", "committed twice", "rejected tokens kept") ) -def test_the_twin_sees_a_broken_protocol(fault, t) -> None: +def test_the_twin_sees_a_broken_protocol(fault) -> None: """The twin comparison is what certifies the commit and the bookkeeping, so it has to fail when either is broken: the partitioned path with the manager's commit left out; a 16-row batch called without work items while the manager still commits (the checkpoint advanced twice); and the small path advanced as if every drafted token had been accepted. Without a fault - the same run passes.""" + the same run passes. At T = 8, the largest token tile.""" path = "small" if fault == "rejected tokens kept" else "all_layer" n = 8 if path == "small" else 16 - bench = _Bench(dict(path=path, n=n, t=t, steps=4), seed=93) + bench = _Bench(dict(path=path, n=n, t=8, steps=4), seed=93) first = _drive(bench, fault) if fault == "none": assert first is None else: - assert first is not None, f"{fault} passed the twin comparison at T={t}" + assert first is not None, f"{fault} passed the twin comparison" @pytest.mark.parametrize("state_dtype", (torch.float32, torch.bfloat16), ids=("fp32", "bf16")) @@ -809,8 +806,6 @@ def test_a_history_off_history_size_is_misread_and_is_refused() -> None: ( ("history", "fewer slots"), ("prev_num_accepted_tokens", "fewer slots"), - ("cache_buf_idx", "fewer slots"), - ("prev_num_accepted_tokens", "strided"), ("cache_buf_idx", "strided"), ), ) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_update.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_update.py index 980c4be80374..8da57911d100 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_update.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_update.py @@ -3,9 +3,9 @@ """GPU test for the fused_recurrent_gated_delta_rule_update entry. The axis this entry exists for is more than one token per sequence per step -- -the speculative verify call -- so the cells sweep `P`, and the two state sinks -(write-back into the pool, per-step states into the rollback buffer) are -driven separately and together. +the speculative verify call -- so the cells are that call as the target makes it +(per-step states into the rollback buffer, the pool untouched) at both pool +dtypes. The write-back sink is driven only by the property and guard tests below. Two tests answer "is this harness blind" rather than a correctness question: the neighbouring `chunk_gated_delta_rule` is wrong on the first indexed diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_sigmoid_gating_delta_rule_update.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_sigmoid_gating_delta_rule_update.py index 4f7a519ce50b..7d0096e9bd80 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_sigmoid_gating_delta_rule_update.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_sigmoid_gating_delta_rule_update.py @@ -2,11 +2,11 @@ # SPDX-License-Identifier: Apache-2.0 """GPU test for the fused_sigmoid_gating_delta_rule_update entry. -Everything here certifies the vendored Triton branch. The target's two decode calls -keep the FlashInfer branch out -- the packed call over an fp32 pool, and the rows form -(no `cu_seqlens`) over a bf16 one -- which `test_certified_branch_is_the_vendored_triton_kernel` -proves. `test_rows_form_is_batch_invariant` shows the property the bf16 call takes the rows -form for. +Everything here certifies the vendored Triton branch. The target's decode call -- the +rows form (no `cu_seqlens`) over a bf16 pool -- and the packed call over an fp32 pool that +the GDN block test uses as its reference keep the FlashInfer branch out, which +`test_certified_branch_is_the_vendored_triton_kernel` proves. +`test_rows_form_is_batch_invariant` shows the property the bf16 call takes the rows form for. """ import zlib @@ -195,9 +195,9 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res def test_certified_branch_is_the_vendored_triton_kernel() -> None: - """Both target decode calls keep the FlashInfer arm out, armed both ways: its entry - point is replaced by one that raises; the packed call over an fp32 pool and the rows - form over a bf16 pool complete, and the packed call over a bf16 pool in the same + """The target's rows form and the reference's fp32 packed call keep the FlashInfer arm + out, armed both ways: its entry point is replaced by one that raises; the packed call + over an fp32 pool and the rows form over a bf16 pool complete, and the packed call over a bf16 pool in the same process does raise -- so the replaced entry point is known to be reachable.""" def _boom(*args, **kwargs): @@ -288,7 +288,7 @@ def test_rows_form_is_the_packed_triton_call_bitwise() -> None: assert torch.equal(rows[9], packed[9]), (pool_dtype, n) -@pytest.mark.parametrize("n", (1, 16, 32)) +@pytest.mark.parametrize("n", (1, 32)) def test_rows_form_cuda_graph_replay_is_bitwise_eager(n) -> None: """The bf16 rows form captured at batch `n` and replayed on new activations, slots and pool content, against the same call run eagerly. The kernel takes the batch as a diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_block.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_block.py index 10c6bdc31fc8..017b2289cfc0 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_block.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_block.py @@ -337,8 +337,9 @@ def warm_triton_chunk(): torch.cuda.synchronize() -@pytest.mark.parametrize("m", (1, 32, 128)) -def test_decode_block(m) -> None: +def test_decode_block() -> None: + """A con-128 decode step; the graph-replay test below runs 32 rows.""" + m = 128 core, L, x = _core(100 + m), _layer(200 + m), _x(m, 300 + m) core_ref = _clone_core(core) st, qsl_long = _decode_state(m) @@ -353,13 +354,12 @@ def test_decode_block(m) -> None: ) -#: (context lengths, generation rows, has_initial_state per context). +#: (context lengths, generation rows, has_initial_state per context): an 8K prefill-only +#: step, and a mixed step of three variable-length contexts (two carrying state) plus +#: generation rows. PREFILL_CASES = ( ((8192,), 0, (False,)), - ((200, 37, 91), 0, (True, False, True)), - ((8065,), 127, (True,)), ((200, 37, 91), 5, (True, False, True)), - ((64,), 1, (False,)), ) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_gated_norm_pdl.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_gated_norm_pdl.py index 7f45e90b392e..3b95ef012e24 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_gated_norm_pdl.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_gated_norm_pdl.py @@ -119,8 +119,8 @@ def _unfused(proj, pool, slots, A_log, dt_bias, w, eps, h, k): ).reshape(rows, HV * V) -@pytest.mark.parametrize("rows", [1, 4]) -def test_matches_the_unfused_composition(rows) -> None: +def test_matches_the_unfused_composition() -> None: + rows = 4 # MAX_ROWS args = _build(dict(rows=rows), seed=900 + rows) pool_unfused = args[1].clone() want = _unfused(args[0], pool_unfused, *args[2:]) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py index b27964b8fd6e..77756f9242aa 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py @@ -6,8 +6,8 @@ then the gated norm entry's) in torch's default band for the bf16 output and the fp32 pool, the pool compared whole so slots the call did not name are gated too. Beyond the cells: agreement with the unfused composition it replaces -(FlashInfer's recurrence, then the gated norm entry) at decode batch sizes, PDL on -and off bitwise, a stale-read check under graph replay (a producer rewrites the +(FlashInfer's recurrence, then the gated norm entry) at the 128-row decode batch, PDL +on and off bitwise, a stale-read check under graph replay (a producer rewrites the projection inside the graph right before the op), and the guards. """ @@ -122,8 +122,8 @@ def _unfused(proj, pool, slots, A_log, dt_bias, w, eps, h, k): ).reshape(rows, HV * V) -@pytest.mark.parametrize("rows", [3, 64, 128]) -def test_matches_the_unfused_composition(rows) -> None: +def test_matches_the_unfused_composition() -> None: + rows = 128 args = _build(dict(rows=rows), seed=900 + rows) pool_unfused = args[1].clone() want = _unfused(args[0], pool_unfused, *args[2:]) @@ -217,10 +217,11 @@ def test_rows_certified_cells(cell, certifying) -> None: # noqa: F811 — fixtu assert (got[1] == 0).all(), "a pad row's output is zero" -@pytest.mark.parametrize("rows", [3, 121]) -def test_rows_then_gated_norm_matches_the_fused_entry(rows) -> None: +def test_rows_then_gated_norm_matches_the_fused_entry() -> None: """The recurrence alone, then the gated norm entry: the fused entry's output within its - band, and the same pool bitwise (one kernel's recurrence).""" + band, and the same pool bitwise (one kernel's recurrence), at a mixed step's 121 + generation rows.""" + rows = 121 args = _build(dict(rows=rows), seed=700 + rows) pool_fused = args[1].clone() fused = op(args[0], pool_fused, *args[2:]) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_in_proj_conv1d_pdl.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_in_proj_conv1d_pdl.py index 6d5b2d1959b7..dcaf4fb3c133 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_in_proj_conv1d_pdl.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_in_proj_conv1d_pdl.py @@ -73,8 +73,8 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 assert torch.equal(pool[slot], state[slot]), f"slot {slot} was written" -@pytest.mark.parametrize("rows", [1, 2]) -def test_matches_the_unfused_composition(rows) -> None: +def test_matches_the_unfused_composition() -> None: + rows = MAX_ROWS x, mat_b, state, taps = _build(rows, 461, seed=5 + rows) slots = torch.tensor([37, 211][:rows], dtype=torch.int32, device="cuda") fused_pool, unfused_pool = state.clone(), state.clone() @@ -94,12 +94,12 @@ def test_matches_the_unfused_composition(rows) -> None: _compare(fused_pool, unfused_pool) -@pytest.mark.parametrize("pdl", [False, True]) -def test_two_steps_under_graph_replay(pdl) -> None: - """Decode runs this under CUDA graphs, layer after layer: capture one step, - replay it twice with new activations. The second replay's conv reads the +def test_two_steps_under_graph_replay() -> None: + """Decode runs this under CUDA graphs with PDL, layer after layer: capture one + step, replay it twice with new activations. The second replay's conv reads the window the first replay advanced, so a stale or misplaced window write fails the second comparison.""" + pdl = True x, mat_b, state, taps = _build(1, 64, seed=3) slots = torch.tensor([9], dtype=torch.int32, device="cuda") pool = state.clone() diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_prefill_conv_post.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_prefill_conv_post.py index dafeb0921e6c..e9ce6b8488e9 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_prefill_conv_post.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_prefill_conv_post.py @@ -2,10 +2,11 @@ # SPDX-License-Identifier: Apache-2.0 """GPU test for the gdn_prefill_conv_post entry. -Every cell is checked twice: against the entry's fp32/fp64 reference, and -against the three-call composition it replaces in the target +Every cell is checked against the entry's fp32/fp64 reference, the conv pool +bitwise. The three-call composition it replaces in the target (`extract_transpose_prefill_slice` -> `causal_conv1d_fwd` -> -`fused_gdn_post_conv`), with the conv pool compared bitwise both times. +`fused_gdn_post_conv`, each certified on its own) is checked once, at the +target-call cell, again with the conv pool compared bitwise. Geometry is a Qwen3.8-Flash-Next linear-attention layer: 16 x 128 key heads, 48 x 128 value heads, qkv_dim 10240, conv width 4, bf16 activations. """ @@ -65,19 +66,17 @@ def _build(spec, seed): g = torch.Generator(device=DEV).manual_seed(seed) lengths = spec["lengths"] n_tok = sum(lengths) - layout = spec.get("layout", "in_proj") def randn(*shape): return torch.randn(*shape, generator=g, device=DEV).to(DTYPE) - if layout == "contiguous": - x, a, b = randn(n_tok, QKV_DIM), randn(n_tok, NUM_V_HEADS), randn(n_tok, NUM_V_HEADS) - else: - proj = randn(n_tok + spec.get("extra_rows", 0), IN_PROJ_WIDTH) - x = proj[:n_tok, :QKV_DIM] - a = proj[:n_tok, A_COL : A_COL + NUM_V_HEADS] - b = proj[:n_tok, A_COL + NUM_V_HEADS :] - assert x.stride(0) == a.stride(0) == IN_PROJ_WIDTH + # "in_proj" and "rows" both carve x, a and b out of one projection; "rows" + # leaves `extra_rows` trailing rows the call must not read. + proj = randn(n_tok + spec.get("extra_rows", 0), IN_PROJ_WIDTH) + x = proj[:n_tok, :QKV_DIM] + a = proj[:n_tok, A_COL : A_COL + NUM_V_HEADS] + b = proj[:n_tok, A_COL + NUM_V_HEADS :] + assert x.stride(0) == a.stride(0) == IN_PROJ_WIDTH weight = (torch.randn(QKV_DIM, CONV_WIDTH, generator=g, device=DEV) * 0.5).to(DTYPE) pool_slots = spec["pool_slots"] # The pool is a per-layer view of a two-layer buffer, as a cache manager @@ -127,30 +126,19 @@ def _unfused(args): return (*out, pool) -def _live_rows(args): - """Token rows of the contexts whose slot is not the pad sentinel.""" - bounds, slots = args[3].tolist(), args[4].tolist() - keep = torch.zeros(args[0].shape[0], dtype=torch.bool, device=DEV) - for s, slot in enumerate(slots): - if slot != PAD: - keep[bounds[s] : bounds[s + 1]] = True - return keep - - -def _compare(out, pool, want, rows=None, what=""): +def _compare(out, pool, want, what=""): q, k, v, g, beta = out q_r, k_r, v_r, g_r, beta_r, pool_r = want for got, ref in zip(out, want[:5]): assert got.dtype == ref.dtype and got.shape == ref.shape and got.is_contiguous() - sel = (lambda t: t[0, rows]) if rows is not None else (lambda t: t) for name, got, ref in (("q", q, q_r), ("k", k, k_r), ("v", v, v_r)): - d = (sel(got).float() - sel(ref).float()).abs() + d = (got.float() - ref.float()).abs() flips = (d > 0).float().mean().item() if d.numel() else 0.0 print( f"{what} {name}: max abs {d.max().item() if d.numel() else 0:.3e}, " f"differing fraction {flips:.2e}" ) - torch.testing.assert_close(sel(got), sel(ref)) + torch.testing.assert_close(got, ref) torch.testing.assert_close(g, g_r, rtol=8 * _APPROX_ULP, atol=G_ATOL) torch.testing.assert_close(beta, beta_r) assert torch.equal(pool, pool_r), "conv pool differs" @@ -162,19 +150,27 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res args = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) pool0 = args[2].clone() want = op.reference(*args) - old = _unfused([*args[:2], pool0.clone(), *args[3:]]) inputs = [t.clone() if isinstance(t, torch.Tensor) else t for t in args] out = op(*args) for i, (arg, before) in enumerate(zip(args, inputs)): if isinstance(arg, torch.Tensor) and i != 2: assert torch.equal(arg, before), f"input {i} must be read-only" _compare(out, args[2], want, what=f"{cell.why[:30]} vs reference") - _compare(out, args[2], old, rows=_live_rows(args), what=f"{cell.why[:30]} vs unfused") if PAD in args[4].tolist(): pad_free = [s for s in range(pool0.shape[0]) if s not in args[4].tolist()] assert torch.equal(args[2][pad_free], pool0[pad_free]), "unnamed slots written" +def test_matches_the_unfused_composition() -> None: + """The target-call cell against the three calls the entry replaces; each of + those is certified on its own, so one representative call is the check.""" + cell = next(c for c in op.CELLS if c.why.startswith("target call")) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) + old = _unfused(args) + out = op(*args) + _compare(out, args[2], old, what=f"{cell.why[:30]} vs unfused") + + def test_carry_and_silu_are_load_bearing() -> None: """The band separates a dropped carry and a dropped SiLU.""" spec = dict(lengths=(5, 70), pool_slots=4, slots=(1, 3), initial=(True, True)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_pack_gdn_decode_qkv.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_pack_gdn_decode_qkv.py index 54fe1ac9e2d0..50a569172aac 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_pack_gdn_decode_qkv.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_pack_gdn_decode_qkv.py @@ -45,83 +45,56 @@ def _compare(out: tuple, ref: tuple) -> None: DEV = "cuda" Q_WIDTH = NUM_K_HEADS * HEAD_K_DIM # 2048 V_WIDTH = NUM_V_HEADS * HEAD_V_DIM # 6144 -#: The `z` half of `in_proj_qkvz`, of which the stock mixer's `mixed_qkv` is a column slice. -Z_WIDTH = NUM_V_HEADS * HEAD_V_DIM def _build(spec, seed): - """`(mixed_qkv, conv_out or None)`: the input, and for a verify batch the + """`(mixed_qkv, conv_out or None)`: the bf16 input, and for a verify batch the `[B, qkv_dim, P]` conv output it was flattened from.""" g = torch.Generator(device=DEV).manual_seed(seed) - dtype = spec.get("dtype", torch.bfloat16) - heads = spec.get("heads", CKPT_HEADS) - qkv_dim = 2 * heads[0] * heads[1] + heads[2] * heads[3] - layout = spec.get("layout", "contiguous") def randn(*shape): - return torch.randn(*shape, generator=g, device=DEV).to(dtype) + return torch.randn(*shape, generator=g, device=DEV).to(torch.bfloat16) if "draft" in spec: b, p = spec["num_decodes"], spec["draft"] - if layout == "conv_dim_major": - conv_out = randn(b, qkv_dim, p) # contiguous; the flatten copies unless b == 1 - else: - # The target: the conv result carries its `[n, C, p]` input view's - # strides, so transposing back and flattening is a contiguous view. - conv_out = randn(b, p, qkv_dim).transpose(1, 2) - mixed = conv_out.transpose(1, 2).reshape(b * p, qkv_dim) - if layout != "conv_dim_major": - assert mixed.is_contiguous() and mixed.data_ptr() == conv_out.data_ptr() + # The target: the conv result carries its `[n, C, p]` input view's + # strides, so transposing back and flattening is a contiguous view. + conv_out = randn(b, p, QKV_DIM).transpose(1, 2) + mixed = conv_out.transpose(1, 2).reshape(b * p, QKV_DIM) + assert mixed.is_contiguous() and mixed.data_ptr() == conv_out.data_ptr() return mixed, conv_out - n = spec["num_tokens"] - if layout == "row_slice_of_projection": - mixed = randn(n + 9, qkv_dim + Z_WIDTH)[9:, :qkv_dim] - assert mixed.stride() == (qkv_dim + Z_WIDTH, 1) - elif layout == "column_stride_2": - mixed = randn(n, 2 * qkv_dim)[:, ::2] - assert mixed.stride() == (2 * qkv_dim, 2) - elif layout == "transposed_view": - mixed = randn(qkv_dim, n).t() - assert mixed.stride() == (1, n) - elif layout == "wide": - mixed = randn(n, qkv_dim + Z_WIDTH) - else: - mixed = randn(n, qkv_dim) - return mixed, None + return randn(spec["num_tokens"], QKV_DIM), None @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name certifying(op) spec = cell.spec - heads = spec.get("heads", CKPT_HEADS) mixed, conv_out = _build(spec, seed=zlib.crc32(cell.why.encode())) snapshot = mixed.clone() - out = op(mixed, *heads) - _compare(out, op.reference(mixed, *heads)) + out = op(mixed, *CKPT_HEADS) + _compare(out, op.reference(mixed, *CKPT_HEADS)) assert all(t.is_contiguous() for t in out) assert torch.equal(mixed, snapshot), "mixed_qkv is read-only" - if mixed.numel(): # an empty allocation has no storage to tell apart - storages = {t.untyped_storage().data_ptr() for t in out} - assert len(storages) == 3 and mixed.untyped_storage().data_ptr() not in storages + storages = {t.untyped_storage().data_ptr() for t in out} + assert len(storages) == 3 and mixed.untyped_storage().data_ptr() not in storages if "draft" in spec: # The verify caller regroups each return as [B, P, heads, dim]: that is # a view of the same storage, and it lands on (sequence, draft) exactly. b, p = spec["num_decodes"], spec["draft"] - num_k_heads, head_k_dim, num_v_heads, head_v_dim = heads grouped = ( - out[0].reshape(b, p, num_k_heads, head_k_dim), - out[1].reshape(b, p, num_k_heads, head_k_dim), - out[2].reshape(b, p, num_v_heads, head_v_dim), + out[0].reshape(b, p, NUM_K_HEADS, HEAD_K_DIM), + out[1].reshape(b, p, NUM_K_HEADS, HEAD_K_DIM), + out[2].reshape(b, p, NUM_V_HEADS, HEAD_V_DIM), ) assert all(gt.data_ptr() == t.data_ptr() for gt, t in zip(grouped, out)) per_token = conv_out.permute(0, 2, 1) # [B, P, qkv_dim], straight off the conv want = ( - per_token.narrow(2, 0, Q_WIDTH).reshape(b, p, num_k_heads, head_k_dim), - per_token.narrow(2, Q_WIDTH, Q_WIDTH).reshape(b, p, num_k_heads, head_k_dim), - per_token.narrow(2, 2 * Q_WIDTH, V_WIDTH).reshape(b, p, num_v_heads, head_v_dim), + per_token.narrow(2, 0, Q_WIDTH).reshape(b, p, NUM_K_HEADS, HEAD_K_DIM), + per_token.narrow(2, Q_WIDTH, Q_WIDTH).reshape(b, p, NUM_K_HEADS, HEAD_K_DIM), + per_token.narrow(2, 2 * Q_WIDTH, V_WIDTH).reshape(b, p, NUM_V_HEADS, HEAD_V_DIM), ) assert all(torch.equal(gt, w) for gt, w in zip(grouped, want)) @@ -163,9 +136,9 @@ def test_bitwise_gate_discriminates_wrong_variants() -> None: 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: +def test_op_rejects_head_dims_that_round_to_different_blocks() -> None: """Loud at compile time, before any launch -- so not an `is_valid` guard.""" + heads = (16, 128, 48, 64) 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"): @@ -192,13 +165,6 @@ def test_a_short_row_reads_the_next_row_and_is_refused() -> None: 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 diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_reset_recurrent_state_rows.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_reset_recurrent_state_rows.py index 84c95ef21b64..55c78cddac33 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_reset_recurrent_state_rows.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_reset_recurrent_state_rows.py @@ -50,7 +50,7 @@ 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_shape, conv_shape = REC_SHAPE, CONV_SHAPE rec_dtype = spec.get("rec_dtype", torch.float32) layout = spec.get("layout", "plain") @@ -78,14 +78,8 @@ def randn(*shape, dtype): 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) + indices = torch.tensor(spec["indices"], dtype=spec.get("index_dtype", torch.int32), device=DEV) + flags = torch.tensor(spec["flags"], dtype=torch.bool, device=DEV) return [rec, indices, flags, conv], backing @@ -117,10 +111,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res 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: @@ -138,42 +128,23 @@ def test_the_following_kernel_sees_the_zeros() -> None: @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.""" + """Why `is_valid` guards only the index tensors: the rest raises itself. One case + per family the callable checks: slot counts, row layout, device.""" 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": + if 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() @@ -217,8 +188,9 @@ def test_guard_hazard_two_d_index_tensors() -> None: _compare(got, want) -@pytest.mark.parametrize("which", ["state_indices", "has_initial_states"]) -@pytest.mark.parametrize("how", ["strided", "two_d"]) +@pytest.mark.parametrize( + "which,how", [("state_indices", "strided"), ("has_initial_states", "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) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_gdn_replay_verify.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_gdn_replay_verify.py index 239bcb1c13f8..cef8c8358a46 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_gdn_replay_verify.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_gdn_replay_verify.py @@ -242,17 +242,10 @@ def _state(gen, cpu, n, t, gating=None): @pytest.mark.parametrize( "n,t,accept", [ - (1, 2, "random"), + # Below the partition threshold: the history wraps in-kernel. (8, 4, "random"), - (CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE, 4, "random"), - (20, 2, "random"), - (18, 3, "random"), - # max_draft_len 5..7: T = 6..8, the replay kernel's ceiling. - (1, 8, "random"), - (8, 8, "random"), + # At the threshold, at max_draft_len 7 (T = 8, the replay kernel's ceiling). (CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE, 8, "random"), - (32, 8, "random"), - (20, 6, "random"), # Every draft accepted: from a fresh slot at T = 8 the second step appends exactly up # to the history's end (`pnat + T == 16`) and every second step after it wraps -- # in-kernel below the partition threshold, through the all-layer commit at and above it. diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_index_share.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_index_share.py index c1c91f33bf60..357e263bc0c7 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_index_share.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_index_share.py @@ -84,43 +84,6 @@ def test_select_is_the_canonical_row_ascending_with_padding_last(): assert row == ids + [-1] * (K - len(ids)), "padding must stay last" -def test_decode_capture_takes_each_request_last_accepted_row(): - p, batch = 4, 3 - accepted = [1, 4, 2] - topk = torch.arange(batch * p * K, dtype=torch.int32, device=DEV).reshape(batch * p, K) - visible = torch.arange(100, 100 + batch * p, dtype=torch.int32, device=DEV) - md = _metadata(batch, accepted) - - _target(m.DecodeTarget)._capture_selection(md, topk, visible, batch, p) - - rows = [s * p + a - 1 for s, a in enumerate(accepted)] - torch.testing.assert_close(md.qsa_shared_topk_indices[SLOT, :batch], topk[rows]) - torch.testing.assert_close(md.qsa_shared_topk_visible_blocks[SLOT, :batch], visible[rows]) - # Other slots and rows past the batch are untouched. - assert torch.all(md.qsa_shared_topk_indices[SLOT, batch:] == -7) - assert torch.all(md.qsa_shared_topk_indices[0] == -7) - assert md.qsa_shared_topk_captured_slots == {SLOT} - - -def test_prefill_capture_takes_chunk_last_row_and_last_accepted_row(): - """Two context requests (chunks of 5 and 3 rows) then two generation - requests (4 rows each, 2 and 4 accepted). The context requests' accepted - count is whatever the worker left there and must not matter.""" - seq_lens = [5, 3, 4, 4] - accepted = [3, 1, 2, 4] - rows_total = sum(seq_lens) - topk = torch.arange(rows_total * K, dtype=torch.int32, device=DEV).reshape(rows_total, K) - visible = torch.arange(200, 200 + rows_total, dtype=torch.int32, device=DEV) - md = _metadata(4, accepted, seq_lens=seq_lens, num_contexts=2) - - _target(m.PrefillTarget)._capture_selection(md, topk, visible) - - rows = [4, 7, 8 + 1, 12 + 3] - torch.testing.assert_close(md.qsa_shared_topk_indices[SLOT, :4], topk[rows]) - torch.testing.assert_close(md.qsa_shared_topk_visible_blocks[SLOT, :4], visible[rows]) - assert md.qsa_shared_topk_captured_slots == {SLOT} - - def test_decode_capture_at_draft_seven_takes_row_s_times_8_plus_a_minus_1(): """max_draft_len 7: eight rows per request, accepted 1..8, so the captured row is `s * 8 + a - 1` -- the first and the last row of a block included.""" @@ -136,25 +99,28 @@ def test_decode_capture_at_draft_seven_takes_row_s_times_8_plus_a_minus_1(): assert rows == [0, 15, 20, 26] torch.testing.assert_close(md.qsa_shared_topk_indices[SLOT, :batch], topk[rows]) torch.testing.assert_close(md.qsa_shared_topk_visible_blocks[SLOT, :batch], visible[rows]) + # Other slots and rows past the batch are untouched. assert torch.all(md.qsa_shared_topk_indices[SLOT, batch:] == -7) + assert torch.all(md.qsa_shared_topk_indices[0] == -7) assert md.qsa_shared_topk_captured_slots == {SLOT} def test_prefill_capture_at_draft_seven(): - """A mixed batch at max_draft_len 7: one 5-row context chunk, then three generation - requests of 8 rows with 8, 1 and 6 accepted.""" - seq_lens = [5, 8, 8, 8] - accepted = [2, 8, 1, 6] + """A mixed batch at max_draft_len 7: two context chunks of 5 and 3 rows, then three + generation requests of 8 rows with 8, 1 and 6 accepted. A context request captures its + chunk's last row whatever accepted count the worker left there.""" + seq_lens = [5, 3, 8, 8, 8] + accepted = [2, 3, 8, 1, 6] rows_total = sum(seq_lens) topk = torch.arange(rows_total * K, dtype=torch.int32, device=DEV).reshape(rows_total, K) visible = torch.arange(400, 400 + rows_total, dtype=torch.int32, device=DEV) - md = _metadata(4, accepted, seq_lens=seq_lens, num_contexts=1) + md = _metadata(5, accepted, seq_lens=seq_lens, num_contexts=2) _target(m.PrefillTarget)._capture_selection(md, topk, visible) - rows = [4, 5 + 7, 13 + 0, 21 + 5] - torch.testing.assert_close(md.qsa_shared_topk_indices[SLOT, :4], topk[rows]) - torch.testing.assert_close(md.qsa_shared_topk_visible_blocks[SLOT, :4], visible[rows]) + rows = [4, 7, 8 + 7, 16 + 0, 24 + 5] + torch.testing.assert_close(md.qsa_shared_topk_indices[SLOT, :5], topk[rows]) + torch.testing.assert_close(md.qsa_shared_topk_visible_blocks[SLOT, :5], visible[rows]) assert md.qsa_shared_topk_captured_slots == {SLOT} diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_kv_range_probe.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_kv_range_probe.py index 57432c1b1199..83cb86883b0e 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_kv_range_probe.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_kv_range_probe.py @@ -42,7 +42,7 @@ def test_a_forward_within_the_e4m3_range_is_silent(monkeypatch): logger.warning.assert_not_called() -@pytest.mark.parametrize("peak", [450.0, 1000.0, float("inf"), float("nan")]) +@pytest.mark.parametrize("peak", [1000.0, float("nan")]) def test_a_magnitude_past_the_e4m3_maximum_is_reported_once(monkeypatch, peak): logger = Mock() monkeypatch.setattr(m, "logger", logger) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_logits_widening.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_logits_widening.py index c369850d3443..c4aea680b2e9 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_logits_widening.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_logits_widening.py @@ -27,7 +27,7 @@ def head(): return lin -@pytest.mark.parametrize("seq_lens", [[1] * 128, [5000, 3000, 1] + [1] * 121, [7]]) +@pytest.mark.parametrize("seq_lens", [[1] * 128, [5000, 3000, 1] + [1] * 121]) def test_matches_the_shell_processor_bitwise(head, seq_lens) -> None: g = torch.Generator(device="cuda").manual_seed(len(seq_lens)) lens = torch.tensor(seq_lens, dtype=torch.int32, device="cuda") diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_spec_rollback.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_spec_rollback.py index 389a0e1efb60..25df5f9bc522 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_spec_rollback.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_spec_rollback.py @@ -261,13 +261,10 @@ def _assert_state(b, expected, what): # compression anchors; one request accepts everything. CASES = { "decode": ((), (5, 30, 61, 100), 4, (1, 2, 4, 3)), - "decode_p2": ((), (0, 31, 64), 2, (1, 2, 1)), "mixed": ((7, 3), (12, 29, 95), 4, (2, 1, 3)), - "mixed_p3": ((40,), (2, 62), 3, (3, 1)), # max_draft_len 7: every acceptance 1..8 once; blocks at 27, 60 and 90 straddle a page, # and an 8-row block completes two compression groups. "decode_p8": ((), (0, 27, 60, 90, 101, 5, 44, 119), 8, (1, 2, 3, 4, 5, 6, 7, 8)), - "mixed_p8": ((7, 3), (12, 29, 95), 8, (8, 1, 5)), } @@ -284,12 +281,12 @@ def test_commit_matches_the_accepted_prefix(case): _assert_state(b, _reference(b), case) -@pytest.mark.parametrize("case", ["mixed", "decode_p8"]) -def test_commit_over_the_strided_planes_of_a_shared_page(case): +def test_commit_over_the_strided_planes_of_a_shared_page(): """The handler stages and restores the index and position planes by advanced indexing, so planes that are strided views of one buffer (an NVFP4 KV cache's side planes) roll back like contiguous ones, and the bytes of the other layers' pages in the same buffer stay as the step wrote them.""" + case = "decode_p8" ctx_lens, gen_starts, p, gen_accepted = CASES[case] b = _Batch( 13, ctx_lens, gen_starts, p, _slots(len(ctx_lens) + len(gen_starts)), shared_page=True @@ -398,7 +395,6 @@ def _single(b, s): #: Per case, the acceptance counts one captured graph is replayed with (the accepted count #: of every generation request, bonus token included). _REPLAYS = { - "decode": [(1, 1, 1, 1), (2, 4, 1, 3), (4, 3, 2, 1), (3, 3, 3, 3)], "decode_p8": [ (1,) * 8, (8, 7, 6, 5, 4, 3, 2, 1), diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_patch_embed_conv3d.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_patch_embed_conv3d.py index 8724f9b5997d..c8d4322ab1c0 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_patch_embed_conv3d.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_patch_embed_conv3d.py @@ -73,48 +73,19 @@ def _operands( return pixel_values, weight, bias_t -def _build(spec, seed): - pixel_values, weight, bias = _operands( - spec["num_patches"], - dtype=spec["dtype"], - seed=seed, - in_channels=spec.get("in_channels", IN_CHANNELS), - kernel=spec.get("kernel", (KT, KH, KW)), - embed_dim=spec.get("embed_dim", EMBED_DIM), - bias=spec.get("bias", True), - ) - layout = spec.get("layout", "flat") - n, row = pixel_values.shape - if layout == "five_d": - pixel_values = pixel_values.view(-1, *weight.shape[1:]) - elif layout == "column_slice": - wide = torch.empty(n, 2 * row, device="cuda", dtype=pixel_values.dtype) - wide[:, :row] = pixel_values - pixel_values = wide[:, :row] - assert not pixel_values.is_contiguous() and pixel_values.stride(-1) == 1 - elif layout == "transposed": - pixel_values = pixel_values.t().contiguous().t() - assert pixel_values.stride(-1) != 1 - return pixel_values, weight, bias - - @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) def test_certified_cells(cell, certifying): # noqa: F811 certifying(op) - pixel_values, weight, bias = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) + spec = cell.spec + pixel_values, weight, bias = _operands( + spec["num_patches"], dtype=spec["dtype"], seed=zlib.crc32(cell.why.encode()) + ) before = pixel_values.clone() - previous = torch.backends.cudnn.allow_tf32 - torch.backends.cudnn.allow_tf32 = False - try: - out = op(pixel_values, weight, bias) - finally: - torch.backends.cudnn.allow_tf32 = previous + out = op(pixel_values, weight, bias) assert out.shape == (pixel_values.shape[0], weight.shape[0]) assert out.dtype == pixel_values.dtype and out.is_contiguous() assert out.data_ptr() != pixel_values.data_ptr() assert torch.equal(pixel_values, before), "the pixels are read-only" - if not pixel_values.is_contiguous() and pixel_values.stride(-1) == 1: - assert torch.equal(out, op(pixel_values.contiguous(), weight, bias)) _compare(out, op.reference(pixel_values, weight, bias)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_pos_embed_interpolate.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_pos_embed_interpolate.py index 78611e4d4c96..deab3b492a53 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_pos_embed_interpolate.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_pos_embed_interpolate.py @@ -7,7 +7,7 @@ torch's own `F.interpolate` in fp64 and shares no arithmetic with the entry, so it gates the structure: which axis is which, how the table is read, and the spatial-merge row order. It sits up to 8.6e-6 from the exact value (its sample -coordinates are float), which is why fp32 cells gate on `op.reference`. +coordinates are float): far inside a bf16 cell's band, though not an fp32 one's. """ import zlib @@ -71,9 +71,8 @@ def test_certified_cells(cell, certifying): # noqa: F811 assert out.dtype == table.dtype assert torch.equal(table, before) and out.data_ptr() != table.data_ptr() torch.testing.assert_close(out, op.reference(table, grid, spec["merge"])) - if spec["dtype"] is not torch.float32: - # The structural gate, from a reference that shares no arithmetic. - torch.testing.assert_close(out, _ref_resize(table, grid, spec["merge"]).to(out.dtype)) + # 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: diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_vision_tower_calls.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_vision_tower_calls.py index d1a339bd5e8a..4096723e8342 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_vision_tower_calls.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_vision_tower_calls.py @@ -124,23 +124,20 @@ def tower(): return M.VisionTower(_core()) -# `(grid, caps)`: frames of 256, 144 and 32 rows, of 400, 96 and 60, of 32 and 8, and of 384. -# Several caps per grid put the frame boundaries at different places, and the short-frame grid -# makes calls whose longest frame is at most 64 rows, where the fused attention kernel is the -# short-frame one. +# `(grid, cap)`: frames of 256, 144 and 32 rows, of 32 and 8, and of 384. The cap cuts the +# mixed grid inside an item and between items, the one-item grid between frames of one video, +# and the short-frame grid into calls whose longest frame is at most 64 rows, where the fused +# attention kernel is the short-frame one. CASES = { - "mixed": ([(2, 16, 16), (3, 12, 12), (6, 4, 8)], (512, 700)), - "image_then_video": ([(1, 20, 20), (4, 8, 12), (1, 6, 10)], (512, 700)), - "short_frames_only": ([(9, 4, 8), (3, 2, 4)], (96, 100)), - "one_item_many_frames": ([(7, 16, 24)], (800, 1200)), + "mixed": ([(2, 16, 16), (3, 12, 12), (6, 4, 8)], 512), + "short_frames_only": ([(9, 4, 8), (3, 2, 4)], 96), + "one_item_many_frames": ([(7, 16, 24)], 800), } -@pytest.mark.parametrize( - "name,cap", [(name, cap) for name, (_, caps) in CASES.items() for cap in caps] -) -def test_calls_of_whole_frames_equal_one_call(tower, monkeypatch, name, cap): - grid = CASES[name][0] +@pytest.mark.parametrize("name", sorted(CASES)) +def test_calls_of_whole_frames_equal_one_call(tower, monkeypatch, name): + grid, cap = CASES[name] pixels = _pixels(grid) grid_thw = torch.tensor(grid) monkeypatch.setattr(M, "_VISION_CALL_MAX_PATCHES", 1 << 20) @@ -156,27 +153,6 @@ def test_calls_of_whole_frames_equal_one_call(tower, monkeypatch, name, cap): _assert_close(chunked, whole) -def test_a_request_that_fits_is_one_call_bit_for_bit(tower, monkeypatch): - grid = CASES["mixed"][0] - pixels = _pixels(grid) - direct = tower._run(pixels, [tuple(row) for row in grid], torch.tensor(grid)) - monkeypatch.setattr(M, "_VISION_CALL_MAX_PATCHES", 1136) - calls = _count_calls(tower, monkeypatch) - out = tower(pixels, torch.tensor(grid)) - assert [rows for rows, _ in calls] == [1136] - assert out.shape == (1136 // 4, OUT) - torch.testing.assert_close(out, direct, rtol=0, atol=0) - - -def test_a_frame_past_the_ceiling_is_refused_before_any_call(tower, monkeypatch): - monkeypatch.setattr(M, "_VISION_CALL_MAX_PATCHES", 256) - calls = _count_calls(tower, monkeypatch) - grid = [(1, 8, 8), (1, 20, 20)] - with pytest.raises(AssertionError, match="one frame is 20x20 = 400 patch rows"): - tower(_pixels(grid), torch.tensor(grid)) - assert calls == [] - - def test_a_request_past_the_old_whole_request_ceiling_runs(tower): # 66 frames of 512x512 pixels: 33 temporal groups of 32x32 patches, 33792 rows in two # calls (32 groups, then 1). The last group's rows, alone in its call, and the first diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_kv_quant.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_kv_quant.py index 6ca70bcdc1c2..53e685592ac8 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_kv_quant.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_kv_quant.py @@ -175,7 +175,9 @@ def test_unquantized_kv_cache_keeps_the_checkpoints_bf16() -> None: core = _core(None) assert core.quant_mode == 0 assert core.kv_pool_dtype == torch.bfloat16 - assert core._pool_geometry(32).kv_dtype == torch.bfloat16 + geometry = core._pool_geometry(32) + assert geometry.kv_dtype == torch.bfloat16 + assert geometry.kv_width == core.head_dim assert core.kv_scale_orig_quant is None and core.kv_scale_quant_orig is None @@ -187,6 +189,7 @@ def test_fp8_kv_cache_stores_e4m3_under_the_fp8_quant_mode() -> None: # planes keep the activation dtype. geometry = core._pool_geometry(32) assert geometry.kv_dtype == torch.float8_e4m3fn + assert geometry.kv_width == core.head_dim assert geometry.dtype == torch.bfloat16 @@ -201,12 +204,6 @@ def test_nvfp4_kv_cache_stores_e2m1_pairs_under_the_nvfp4_quant_mode() -> None: assert geometry.dtype == torch.bfloat16 -@pytest.mark.parametrize("algo", [None, QuantAlgo.FP8]) -def test_unquantized_and_fp8_pools_are_a_full_head_wide(algo) -> None: - core = _core(algo) - assert core._pool_geometry(32).kv_width == core.head_dim - - @pytest.mark.parametrize("algo", [QuantAlgo.INT8, QuantAlgo.FP8_BLOCK_SCALES]) def test_other_kv_cache_algorithms_are_refused(algo) -> None: """The QSA backend serves its sparse path over half, bf16, fp8 and NVFP4 pools and falls diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_moe_scalars.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_moe_scalars.py index a976733d0957..35282a054b47 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_moe_scalars.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_moe_scalars.py @@ -47,7 +47,7 @@ def test_uniform_isc2_keeps_the_per_expert_products_bitwise() -> None: assert torch.equal(got["alpha1"], w["l0_e_isc1"] * w["l0_e_ws2_1"]) -@pytest.mark.parametrize("spread", [1.0, 1.0001, 3.0, 40.0]) +@pytest.mark.parametrize("spread", [1.0001, 40.0]) def test_isc2_spread_is_folded_into_alpha2(spread: float) -> None: """Whatever the spread, `global_sf * alpha2[e]` is expert e's own `ws2_2[e]`, the product trtllm-gen forms from `(1 / isc2[e]) * (isc2[e] * ws2_2[e])`.""" diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_multimodal_embeds.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_multimodal_embeds.py index 48d869234d32..268f9a6ea448 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_multimodal_embeds.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_multimodal_embeds.py @@ -264,9 +264,13 @@ def test_the_tower_calls_whole_frames_and_concatenates_their_rows_in_order(monke def test_a_request_under_the_ceiling_is_one_call(monkeypatch): grid = [(2, 4, 4), (1, 4, 8)] tower = _tower(monkeypatch, 64) - tower(_pixels(grid), torch.tensor(grid)) + pixels = _pixels(grid) + out = tower(pixels, torch.tensor(grid)) assert len(tower.calls) == 1 assert tower.calls[0][2].tolist() == [list(row) for row in grid] + # The one call sees the whole request and its output is the request's, row for row. + assert torch.equal(tower.calls[0][0], pixels) + assert torch.equal(out, pixels[:: MERGE * MERGE, :1].expand(-1, HIDDEN)) def test_a_request_past_the_old_whole_request_ceiling_is_served(monkeypatch): diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_pool_contract.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_pool_contract.py index 93e0bd76cc29..77909197a72d 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_pool_contract.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_pool_contract.py @@ -84,7 +84,6 @@ def test_valid_pools_pass() -> None: ) nvfp4_kv = _t(9, 2, G.heads_kv, G.tokens_per_block, 128, dtype=torch.uint8) M.check_attention_pools(G_NVFP4, 3, _t(9, 32, 1, 128), nvfp4_kv) - M.check_kv_scale_planes(G_NVFP4, 3, nvfp4_kv, _t(9, 2, 2, 32, 16, dtype=torch.uint8)) for dtype in (torch.float32, torch.bfloat16): M.check_recurrent_pools(G, _conv(), _ssm(dtype)) M.check_ple_pools( @@ -120,7 +119,6 @@ def test_valid_pools_pass() -> None: [ (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"), ], ) @@ -205,8 +203,6 @@ def test_kv_scale_planes_must_hold_the_data_pools_pages_per_slot() -> None: (None, "did not allocate"), (_t(9, 2, 2, 32, 16), "torch.uint8"), (_t(9, 2, 2, 32, 8, dtype=torch.uint8), "NVFP4 block scales"), - (_t(9, 2, 4, 32, 16, dtype=torch.uint8), "NVFP4 block scales"), - (_t(9, 3, 2, 32, 16, dtype=torch.uint8), "same slots"), (_t(8, 2, 2, 32, 16, dtype=torch.uint8), "same slots"), ], ) @@ -347,17 +343,13 @@ def test_shared_selection_rejects_missing_row_and_width() -> None: M.check_shared_selection(shared, {3: 1}, 3, 256) -def test_routing_admits_exactly_the_moe_backends_the_target_implements() -> None: - """`route` declines a backend request the constructor would assert on.""" +def test_routing_admits_exactly_what_the_target_implements() -> None: + """`route` declines a MoE backend, a draft length or a recurrent-state dtype the + constructor or the first bind would assert on. `test_modeling_v2_routing.py` pins + routing's behavior at today's values; this keeps routing's constants and the target's + from drifting apart.""" assert set(routing._MOE_BACKENDS) == set(M._TRUNK_MOE_BACKENDS) - - -def test_routing_admits_exactly_the_draft_lengths_the_target_implements() -> None: - """`route` declines a draft length the constructor or the first bind would assert on.""" assert routing._MAX_DRAFT_LEN == M._MAX_VERIFY_WIDTH - 1 assert routing._MAX_LEGACY_DRAFT_LEN == M._MAX_LEGACY_VERIFY_WIDTH - 1 assert routing._MAX_CUTEDSL_DRAFT_LEN == M._CUTEDSL_MAX_DRAFT_LEN - - -def test_routing_admits_exactly_the_state_dtypes_the_target_serves() -> None: assert routing._SSM_STATE_DTYPES == M._SSM_STATE_DTYPES diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_validate_request.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_validate_request.py index 396a99268b1c..b4ee457febf5 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_validate_request.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_validate_request.py @@ -136,9 +136,7 @@ def test_the_frame_bound_is_per_frame_not_per_item(): _admit(_request(video=[(4, 128, 128), (1, 192, 192)])) -@pytest.mark.parametrize( - "grid", [[(0, 4, 4)], [(1, 4, 5)], [(1, 3, 4)], [(1, 0, 4)]], ids=lambda g: str(g) -) +@pytest.mark.parametrize("grid", [[(0, 4, 4)], [(1, 4, 5)]], ids=lambda g: str(g)) def test_a_malformed_grid_is_refused(grid): request = _request(image=[(1, 4, 4)]) request.py_multimodal_data["image"]["image_grid_thw"] = torch.tensor(grid) @@ -175,7 +173,7 @@ def test_missing_position_ids_fail_the_request(): _admit(request) -@pytest.mark.parametrize("shape", [(3, 1, 4), (3, 1, 40), (2, 1, 12 + TEXT_TOKENS), (3, 12 + 5)]) +@pytest.mark.parametrize("shape", [(3, 1, 40), (2, 1, 12 + TEXT_TOKENS), (3, 12 + 5)]) def test_mis_sized_position_ids_fail_the_request(shape): request = _request(video=VIDEOS) request.py_multimodal_data["mrope_config"]["mrope_position_ids"] = torch.zeros(*shape) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_weight_table.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_weight_table.py index 8eea4907b6c4..7282046fbe35 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_weight_table.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_weight_table.py @@ -27,6 +27,8 @@ weights as W, ) +pytestmark = pytest.mark.cpu_only + # Qwen3.8-Flash-Next's own numbers, so the comparison is against the shapes this target # actually ships rather than a toy configuration. _TEXT = dict( @@ -382,7 +384,6 @@ def test_cute_dsl_expert_copy_rows_only_when_kept(backend, prefill_copy): @pytest.mark.parametrize( "backend, max_num_tokens, env, kept", [ - ("TRTLLM", 16384, None, True), ("TRTLLM", 4096, None, True), ("TRTLLM", 4095, None, False), ("TRTLLM", 16384, "0", False), From 8e12b6606e7da9b73946b828c80b1d3349af9930 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Wed, 21 Oct 2026 02:31:18 -0700 Subject: [PATCH 25/28] [TRTLLM-17033][test] Qwen3.8-Flash-Next accuracy: six cases, each engine running every check that fits One MTP1 engine runs gsm8k with its acceptance floor, MMMU and the videos; one CuteDSL + fp8 KV + bf16 state MTP3 engine runs gsm8k and the pass keys; MTP7 pass keys move to NVFP4 KV over the bf16 state. The forced-acceptance case is dropped. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../accuracy/modeling_v2/test_qwen4_exp.py | 473 ++++++++---------- .../references/acceptance_length.yaml | 40 +- .../test_lists/test-db/l0_b300.yml | 33 +- 3 files changed, 228 insertions(+), 318 deletions(-) diff --git a/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py b/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py index 45aa42834bea..949051ff1800 100644 --- a/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py +++ b/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py @@ -129,8 +129,6 @@ def _acceptance_length_from_requests(outputs, draft_len: int) -> float: _PASSKEY_LENGTHS = (3000, 4000, 5000, 6000) _PASSKEY_LENGTHS_LONG = (3000, 6000, 9000, 12000) # 33 prompts, each past the indexer's budget, 32 of which one batch carries. -_FORCED_LENGTHS = tuple(range(2200, 6400, 400)) -_FORCED_TOKENS = 64 _PASSKEY_DEPTHS = (0.15, 0.5, 0.85) _PASSKEY_MAX_TOKENS = 160 @@ -364,32 +362,28 @@ class TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1(LlmapiAccuracyTestHarness): GSM8K_MAX_OUTPUT_LEN = 512 # `mamba_ssm_cache_dtype` is float32 because the checkpoint declares - # `mamba_ssm_dtype: float32` (the engine default is bfloat16); the bf16-state cases - # set it to bfloat16. `dtype` is left at the checkpoint's bf16 KV, which is what the - # anchors were measured at; only the fp8 and NVFP4 KV cases set it. + # `mamba_ssm_dtype: float32` (the engine default is bfloat16); the NVFP4- and fp8-KV cases + # set it to bfloat16, what `auto` resolves to for this checkpoint. `dtype` is left at the + # checkpoint's bf16 KV, which is what the anchors were measured at; only the fp8 and NVFP4 + # KV cases set it. KV_CACHE = KvCacheConfig( free_gpu_memory_fraction=0.6, enable_block_reuse=True, mamba_ssm_cache_dtype="float32", ) - KV_CACHE_BF16_STATE = KvCacheConfig( - free_gpu_memory_fraction=0.6, - enable_block_reuse=True, - mamba_ssm_cache_dtype="bfloat16", - ) - KV_CACHE_FP8 = KvCacheConfig( + KV_CACHE_FP8_BF16_STATE = KvCacheConfig( dtype="fp8", free_gpu_memory_fraction=0.6, enable_block_reuse=True, - mamba_ssm_cache_dtype="float32", + mamba_ssm_cache_dtype="bfloat16", ) # Block reuse is on, which (with chunked prefill and the MTP drafts) is what lets the # attention op read an NVFP4 pool: its kernels exist for the paged-context path only. - KV_CACHE_NVFP4 = KvCacheConfig( + KV_CACHE_NVFP4_BF16_STATE = KvCacheConfig( dtype="nvfp4", free_gpu_memory_fraction=0.6, enable_block_reuse=True, - mamba_ssm_cache_dtype="float32", + mamba_ssm_cache_dtype="bfloat16", ) # Block reuse is off for the vision path, which requires it. KV_CACHE_MULTIMODAL = KvCacheConfig( @@ -429,39 +423,114 @@ class TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1(LlmapiAccuracyTestHarness): # Acceptance-length anchors in references/acceptance_length.yaml, keyed by # target and variant. ACCEPTANCE_KEY = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp1" - ACCEPTANCE_KEY_MTP3_CUTEDSL_BF16_STATE = ( - "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_cutedsl_bf16_state" - ) - ACCEPTANCE_KEY_MTP7_BF16_STATE_LONG_CONTEXT = ( - "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp7_bf16_state_long_context" - ) ACCEPTANCE_KEY_LONG_CONTEXT = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_long_context" - ACCEPTANCE_KEY_FP8_KV_LONG_CONTEXT = ( - "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_fp8_kv_long_context" + ACCEPTANCE_KEY_OPT_IN = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_cutedsl_fp8_kv_bf16_state" + ACCEPTANCE_KEY_OPT_IN_LONG_CONTEXT = ( + "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_cutedsl_fp8_kv_bf16_state_long_context" ) - ACCEPTANCE_KEY_NVFP4_KV_LONG_CONTEXT = ( - "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_nvfp4_kv_long_context" + ACCEPTANCE_KEY_MTP7_LONG_CONTEXT = ( + "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp7_nvfp4_kv_bf16_state_long_context" ) - def _gsm8k_with_acceptance(self, mocker, monkeypatch, engine, speculative_config, key, label): - """gsm8k against the reference row the engine resolves, then the acceptance length. + def _gsm8k_and_acceptance(self, mocker, llm, key, label): + """gsm8k on ``llm``, then the same run's acceptance length against ``key``. - The acceptance length of the same run is held to the floor under ``key``. + gsm8k is read against the reference row the engine resolves; the acceptance length comes + from the engine's iteration statistics, so the engine needs ``enable_iter_perf_stats``. """ mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", self.GSM8K_MAX_OUTPUT_LEN) + GSM8K(self.MODEL_NAME).evaluate(llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) + acceptance_length = compute_acceptance_length(llm) + print(f"[AL] {label} acceptance_length = {acceptance_length:.3f}") + assert_acceptance_length(key, acceptance_length) - with LLM( - self.MODEL_PATH, - speculative_config=speculative_config, - enable_iter_perf_stats=True, - **engine, - **_require_target(monkeypatch, **_VALIDATE), - ) as llm: - task = GSM8K(self.MODEL_NAME) - task.evaluate(llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) - acceptance_length = compute_acceptance_length(llm) - print(f"[AL] {label} acceptance_length = {acceptance_length:.3f}") - assert_acceptance_length(key, acceptance_length) + def _passkeys(self, llm, lengths, draft_len, key, label): + """Pass key retrieval on ``llm`` over prompts past the indexer's budget. + + Every key must be stated back, and two identical requests far apart in one batch + (different slots) must give the same tokens, which fails if a request's selection leaks + into another slot. The acceptance length, from the requests' own metrics, is held to the + floor under ``key``. + """ + tokenizer = AutoTokenizer.from_pretrained(self.MODEL_PATH) + prompts = _passkey_prompts(tokenizer, lengths) + texts = [prompt for prompt, _ in prompts] + [prompts[len(prompts) // 2][0]] + outputs = llm.generate( + texts, + SamplingParams(max_tokens=_PASSKEY_MAX_TOKENS, temperature=0, return_perf_metrics=True), + use_tqdm=False, + ) + missed = [ + k + for (_, k), output in zip(prompts, outputs) + if k not in output.outputs[0].text[:_PASSKEY_MAX_TOKENS] + ] + assert not missed, f"pass keys not stated back: {missed} of {[k for _, k in prompts]}" + twins = outputs[len(prompts) // 2], outputs[-1] + assert twins[0].outputs[0].token_ids == twins[1].outputs[0].token_ids + acceptance_length = _acceptance_length_from_requests(outputs, draft_len) + print(f"[AL] {label} acceptance_length = {acceptance_length:.3f}") + assert_acceptance_length(key, acceptance_length) + + def _mmmu(self, llm): + """MMMU on ``llm`` against the reference row the engine resolves.""" + # MMMU adds <|endoftext|> to the stop token, as the built-in model's + # MMMU test does. + sampling_params = SamplingParams( + max_tokens=MMMU.MAX_OUTPUT_LEN, + truncate_prompt_tokens=MMMU.MAX_INPUT_LEN, + stop="<|endoftext|>", + ) + MMMU(self.MODEL_NAME).evaluate( + llm, + sampling_params=sampling_params, + extra_evaluator_kwargs=dict(chat_template_kwargs=dict(enable_thinking=False)), + ) + + def _video_checks(self, llm): + """Synthetic videos through ``llm``. + + Three videos of a moving square (frames built in memory, since the container has no video + decoder) must each be answered with the right colour and direction. A request mixing an + image and a video, in either order, must reach the forward: its ``mm_item_order`` + manifest has to survive the input processor, the request and the admission check, which + refuses a mixed request without a usable one. A frame past the vision tower's patch + bound (a 3072x3072 image is 192x192 = 36864 patches) must fail as that request's own + error, and the request after it must still complete; ``max_seq_len`` is raised so that + image reaches the admission check rather than the sequence length check. The order of + the splice itself is pinned by ``test_modeling_v2_qwen4_exp_multimodal_embeds``. + """ + tokenizer = AutoTokenizer.from_pretrained(self.MODEL_PATH) + processor = AutoProcessor.from_pretrained( + self.MODEL_PATH, use_fast=True, trust_remote_code=True + ) + greedy = SamplingParams(max_tokens=96, temperature=0) + + def ask(text, items): + request = _media_request(tokenizer, processor, text, items) + return llm.generate_async(request, sampling_params=greedy).result().outputs[0].text + + for name, (_, _, colour, directions) in _VIDEOS.items(): + answer = ask(_VIDEO_QUESTION, [("video", _synthetic_video(name))]).lower() + assert colour in answer and any(d in answer for d in directions), (name, answer) + + # Each label carries its own item's colour. + image = _synthetic_image((250, 200, 0)) + for items, video_colour in ( + ([("video", _synthetic_video("blue_leftward")), ("image", image)], "blue"), + ([("image", image), ("video", _synthetic_video("red_rightward"))], "red"), + ): + answer = ask(_MIXED_QUESTION, items).lower() + assert re.search(r"image:\W*yellow", answer), answer + assert re.search(rf"video:\W*{video_colour}", answer), answer + + with pytest.raises(RequestError, match="frames of 192x192 = 36864 patches"): + ask( + "What shape and colour is in the image?", + [("image", _synthetic_image(_RED, 3072))], + ) + answer = ask(_VIDEO_QUESTION, [("video", _synthetic_video("red_rightward"))]).lower() + assert "red" in answer and "right" in answer, answer @skip_not_sm103 def test_gsm8k(self, mocker, monkeypatch): @@ -477,52 +546,75 @@ def test_gsm8k(self, mocker, monkeypatch): task.evaluate(llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) @skip_not_sm103 - def test_mtp1_gsm8k_and_acceptance(self, mocker, monkeypatch): - """MTP with one draft: gsm8k against the MTP reference, plus acceptance length. - - A miscomputed draft layer only makes generation slower (every draft is - rejected, the text stays correct), so acceptance length is the only - detector. The floor in references/acceptance_length.yaml is a collapse - tripwire. + def test_mtp1_gsm8k_mmmu_and_video(self, mocker, monkeypatch): + """gsm8k with its acceptance-length floor, MMMU and synthetic videos on one MTP engine. + + - gsm8k against the MTP row of references/gsm8k.yaml. A miscomputed draft layer only + makes generation slower (every draft rejected, the text still correct), so the + acceptance length of the same run is the only detector; its floor is a collapse + tripwire. + - MMMU against the MTP row of references/mmmu.yaml: the vision tower and the mrope path. + - Videos (`_video_checks`): the only case that sends a video. + + Block reuse is off, as the vision path requires; it does not change the text numerics. + ``max_seq_len`` is raised for the video checks' oversized image. """ - self._gsm8k_with_acceptance( - mocker, monkeypatch, self.ENGINE, self.MTP1, self.ACCEPTANCE_KEY, "mtp1" - ) + engine = dict(self.ENGINE, kv_cache_config=self.KV_CACHE_MULTIMODAL, max_seq_len=16384) + with LLM( + self.MODEL_PATH, + speculative_config=self.MTP1, + enable_iter_perf_stats=True, + **engine, + **_require_target(monkeypatch, **_VALIDATE), + ) as llm: + self._gsm8k_and_acceptance(mocker, llm, self.ACCEPTANCE_KEY, "mtp1") + self._mmmu(llm) + self._video_checks(llm) @skip_not_sm103 - def test_mtp3_cutedsl_bf16_ssm_state_gsm8k_and_acceptance(self, mocker, monkeypatch): - """Three drafts over a CuteDSL trunk and a bf16 recurrent state: the MTP row and a floor. - - The opt-in storage and kernel choices in one engine: the trunk's routed experts run on - CuteDSL next to the draft layer's FP8 block-scale experts, and the recurrent state is - bf16 (what ``auto`` resolves to for this checkpoint), rounded at every decode token, - prefill call and cached-replay commit. Three drafts give the verify step four rows per - request on every attention and recurrent-state call. Acceptance length detects a trunk - hidden state or a drifted state the draft layer cannot use, which accuracy does not. + def test_mtp3_cutedsl_fp8_kv_bf16_state(self, mocker, monkeypatch): + """gsm8k and 3k-6k-token pass keys with three drafts over the opt-in choices. + + One engine runs gsm8k with an acceptance-length floor, then the pass keys. The trunk's + routed experts run on CuteDSL next to the draft layer's FP8 block-scale experts; the + KV cache is FP8 (unit scales) and the recurrent state bf16 (what ``auto`` + resolves to for this checkpoint), rounded at every decode token, prefill call and + cached-replay commit. gsm8k reads the FP8-KV MTP row of references/gsm8k.yaml; its + prompts stay on the dense attention path. The pass keys run the FP8 sparse store and + GQA kernels of the target and of the draft layer (IndexShare off, so the draft layer + selects on every step). Three drafts give the verify step four rows per request. + Acceptance length detects a trunk hidden state or a drifted state the draft layer + cannot use, which accuracy does not. """ engine = dict( self.ENGINE, moe_config=MoeConfig(backend="CUTEDSL"), - kv_cache_config=self.KV_CACHE_BF16_STATE, - ) - self._gsm8k_with_acceptance( - mocker, - monkeypatch, - engine, - self.MTP3, - self.ACCEPTANCE_KEY_MTP3_CUTEDSL_BF16_STATE, - "mtp3 cutedsl bf16 state", + kv_cache_config=self.KV_CACHE_FP8_BF16_STATE, + sparse_attention_config=QSASparseAttentionConfig(index_share_for_mtp_iteration=False), ) + with LLM( + self.MODEL_PATH, + speculative_config=self.MTP3, + enable_iter_perf_stats=True, + **engine, + **_require_target(monkeypatch, **_VALIDATE), + ) as llm: + self._gsm8k_and_acceptance(mocker, llm, self.ACCEPTANCE_KEY_OPT_IN, "mtp3 opt-in") + self._passkeys( + llm, + _PASSKEY_LENGTHS, + self.MTP3.max_draft_len, + self.ACCEPTANCE_KEY_OPT_IN_LONG_CONTEXT, + "mtp3 opt-in long context", + ) def _long_context_passkey( - self, monkeypatch, speculative_config, engine, index_share, lengths, key_name, label + self, monkeypatch, speculative_config, engine, index_share, lengths, key, label ): - """Pass key retrieval over prompts past the indexer's budget, at ``speculative_config``. + """`_passkeys` with the executor in this process, checking the op calls too. - Runs with the executor in this process to read the draft layer's op calls; the - acceptance length comes from the requests' own metrics. Two identical requests in one - batch must produce the same tokens, which fails if a request's shared selection leaks - into another slot. + The draft layer's IndexShare reuse op must run exactly when IndexShare is on, and no + prefill graph may replay. """ monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") reuse_calls = [] @@ -533,18 +625,13 @@ def counting_reuse(*args, **kwargs): return reuse(*args, **kwargs) monkeypatch.setattr(target_modeling, "qsa_mtp_reuse_topk", counting_reuse) + replays = _count_prefill_graph_replays(monkeypatch) engine = dict( engine, sparse_attention_config=QSASparseAttentionConfig( index_share_for_mtp_iteration=index_share ), ) - tokenizer = AutoTokenizer.from_pretrained(self.MODEL_PATH) - prompts = _passkey_prompts(tokenizer, lengths) - replays = _count_prefill_graph_replays(monkeypatch) - # One prompt twice, far apart in the batch, so the two copies hold different slots. - texts = [prompt for prompt, _ in prompts] + [prompts[len(prompts) // 2][0]] - with LLM( self.MODEL_PATH, speculative_config=speculative_config, @@ -552,66 +639,33 @@ def counting_reuse(*args, **kwargs): **_require_target(monkeypatch, **_VALIDATE), ) as llm: assert _target_core(llm).mtp_index_share == index_share - outputs = llm.generate( - texts, - SamplingParams( - max_tokens=_PASSKEY_MAX_TOKENS, temperature=0, return_perf_metrics=True - ), - use_tqdm=False, - ) - - missed = [ - key - for (_, key), output in zip(prompts, outputs) - if key not in output.outputs[0].text[:_PASSKEY_MAX_TOKENS] - ] - assert not missed, f"pass keys not stated back: {missed} of {[key for _, key in prompts]}" - twins = outputs[len(prompts) // 2], outputs[-1] - assert twins[0].outputs[0].token_ids == twins[1].outputs[0].token_ids + self._passkeys(llm, lengths, speculative_config.max_draft_len, key, label) assert not replays, replays if index_share: assert reuse_calls, "IndexShare is on and the reuse op never ran" else: assert not reuse_calls, "IndexShare is off and the reuse op ran" - acceptance_length = _acceptance_length_from_requests( - outputs, speculative_config.max_draft_len - ) - print(f"[AL] {label} acceptance_length = {acceptance_length:.3f}") - assert_acceptance_length(getattr(self, key_name), acceptance_length) - - @pytest.mark.parametrize( - "kv_cache, index_share, key_name", - [ - pytest.param("KV_CACHE", True, "ACCEPTANCE_KEY_LONG_CONTEXT", id="index_share"), - pytest.param("KV_CACHE_FP8", False, "ACCEPTANCE_KEY_FP8_KV_LONG_CONTEXT", id="fp8_kv"), - pytest.param( - "KV_CACHE_NVFP4", False, "ACCEPTANCE_KEY_NVFP4_KV_LONG_CONTEXT", id="nvfp4_kv" - ), - ], - ) + @skip_not_sm103 - def test_mtp3_long_context_passkey(self, monkeypatch, kv_cache, index_share, key_name): - """Pass key retrieval over 3k-6k-token prompts, where the sparse attention path runs. - - The gsm8k cases never leave the dense path (their prompts are under the indexer's - 2048-token budget), so this is the only case that runs the target's QSA selection, the - draft layer's selection and, over an fp8 or NVFP4 pool, the sparse store and GQA kernels - of that format (the draft layer's included) end to end. A pass key stated back - correctly needs the right blocks selected at every depth. - - ``index_share`` makes the draft loop reuse the selection of its first step; the - acceptance length then must not fall below the floor taken from the same prompts - with the draft layer selecting on every step, and the reuse op must have run. + def test_mtp3_long_context_passkey(self, monkeypatch): + """Pass key retrieval over 3k-6k-token prompts with three drafts and IndexShare. + + The KV cache is bf16; the prompts are long enough to run the sparse attention path. The + gsm8k prompts never leave the dense path (they are under the indexer's 2048-token + budget), so the pass-key cases are the ones that run the target's QSA selection and the + draft layer's end to end; a pass key stated back correctly needs the right blocks + selected at every depth. IndexShare makes the draft loop reuse the selection of its + first step: the acceptance length must not fall below the floor taken from the same + prompts with the draft layer selecting on every step, and the reuse op must have run. """ - engine = dict(self.ENGINE, kv_cache_config=getattr(self, kv_cache)) self._long_context_passkey( monkeypatch, self.MTP3, - engine, - index_share, + self.ENGINE, + True, _PASSKEY_LENGTHS, - key_name, - f"mtp3 long context {kv_cache}", + self.ACCEPTANCE_KEY_LONG_CONTEXT, + "mtp3 long context", ) @skip_not_sm103 @@ -708,158 +762,27 @@ def run(engine): @skip_not_sm103 def test_mtp7_long_context_passkey(self, monkeypatch): - """Pass key retrieval over 3k-12k-token prompts: seven drafts, IndexShare, bf16 state. - - Seven drafts run the sparse path with eight rows per request in the verify step, the - cached replay kernel's ceiling (``T = 8``), and draft decode steps that reuse the first - step's selection up to six positions past it, where the groups completed since the - capture displace two of its ids. The recurrent state is bf16, what ``auto`` resolves to - for this checkpoint: it is rounded at every cached-replay checkpoint commit and once - per prefill chunk, so the rounding adds up with the length of the prompt and of the - answer, and pass keys stated back over 3k-12k tokens bound that drift. 12k tokens past - the 8k default needs ``max_seq_len`` 16384. + """Pass key retrieval over 3k-12k-token prompts with seven drafts and IndexShare. + + The KV cache is NVFP4 and the recurrent state bf16. Seven drafts run the sparse path + with eight rows per request in the verify step, the cached replay kernel's ceiling + (``T = 8``), and draft decode steps that reuse the first step's selection up to six + positions past it, where the groups completed since the capture displace two of its + ids. The NVFP4 pool runs the target's and the draft layer's NVFP4 sparse store and GQA + kernels, and the context rows' dequantized scratch. The bf16 state is rounded at every + cached-replay checkpoint commit and once per prefill chunk, so the rounding adds up + with the length of the prompt and of the answer; pass keys stated back over 3k-12k + tokens bound that drift. 12k tokens past the 8k default needs ``max_seq_len`` 16384. """ - engine = dict(self.ENGINE, max_seq_len=16384, kv_cache_config=self.KV_CACHE_BF16_STATE) + engine = dict( + self.ENGINE, max_seq_len=16384, kv_cache_config=self.KV_CACHE_NVFP4_BF16_STATE + ) self._long_context_passkey( monkeypatch, self.MTP7, engine, True, _PASSKEY_LENGTHS_LONG, - "ACCEPTANCE_KEY_MTP7_BF16_STATE_LONG_CONTEXT", - "mtp7 bf16 state long context", - ) - - @pytest.mark.parametrize( - "forced, accepted_rows", - [pytest.param("1e-9", 1, id="all_rejected")], - ) - @skip_not_sm103 - def test_mtp7_forced_acceptance_extremes(self, monkeypatch, forced, accepted_rows): - """Seven drafts with every draft rejected, at batch 32. - - ``TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS`` overrides the sampler's acceptance: a - vanishing fraction accepts the target's own token and nothing else. 32 requests over - 2k-6k-token prompts run the sparse path, the all-layer replay commit and the CUDA graph - at 32 x 8 = 256 verify rows with the most rows rolled back at every step (seven). The - acceptance length, read from the requests' own metrics, must be exactly the forced row - count: the acceptance bookkeeping of the verify step moved by the number of rows - forced, not by what the drafts say. - - The text under forced acceptance is not compared with anything: the drafts a forced - step accepts are not the target's tokens, and a run's tokens depend on the verify - step's row count through the GEMM and attention kernel selection, so a run with - another draft length or none differs from this one even with every draft rejected. - The rolled-back state itself is pinned bitwise by ``test_modeling_v2_spec_rollback``. - """ - tokenizer = AutoTokenizer.from_pretrained(self.MODEL_PATH) - prompts = [prompt for prompt, _ in _passkey_prompts(tokenizer, _FORCED_LENGTHS)][:32] - env = {"TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS": forced} - - with LLM( - self.MODEL_PATH, - speculative_config=self.MTP7, - **self.ENGINE, - **_require_target(monkeypatch, **env, **_VALIDATE), - ) as llm: - outputs = llm.generate( - prompts, - SamplingParams( - max_tokens=_FORCED_TOKENS, - temperature=0, - ignore_eos=True, - return_perf_metrics=True, - ), - use_tqdm=False, - ) - - assert [len(o.outputs[0].token_ids) for o in outputs] == [_FORCED_TOKENS] * len(prompts) - acceptance_length = _acceptance_length_from_requests(outputs, self.MTP7.max_draft_len) - print(f"[AL] mtp7 forced acceptance of {accepted_rows} rows: {acceptance_length:.3f}") - assert acceptance_length == pytest.approx(accepted_rows, abs=0.05) - - @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 text cases above) because the vision path - requires it. - """ - self._mmmu(monkeypatch, self.ENGINE) - - def _mmmu(self, monkeypatch, engine): - engine = dict(engine, kv_cache_config=self.KV_CACHE_MULTIMODAL) - - # 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|>", + self.ACCEPTANCE_KEY_MTP7_LONG_CONTEXT, + "mtp7 nvfp4 kv bf16 state long context", ) - - with LLM( - self.MODEL_PATH, - speculative_config=self.MTP1, - **engine, - **_require_target(monkeypatch, **_VALIDATE), - ) as llm: - task = MMMU(self.MODEL_NAME) - task.evaluate( - llm, - sampling_params=sampling_params, - extra_evaluator_kwargs=dict(chat_template_kwargs=dict(enable_thinking=False)), - ) - - @skip_not_sm103 - def test_mtp1_video(self, monkeypatch): - """Synthetic videos through the engine: the only case that sends a video. - - Three videos of a moving square (frames built in memory, since the container has no video - decoder) must each be answered with the right colour and direction. A request mixing an - image and a video, in either order, must reach the forward: its ``mm_item_order`` - manifest has to survive the input processor, the request and the admission check, which - refuses a mixed request without a usable one. A frame past the vision tower's patch - bound (a 3072x3072 image is 192x192 = 36864 patches) must fail as that request's own - error, and the request after it must still complete; ``max_seq_len`` is raised so that - image reaches the admission check rather than the sequence length check. The order of - the splice itself is pinned by ``test_modeling_v2_qwen4_exp_multimodal_embeds``. - """ - engine = dict(self.ENGINE, kv_cache_config=self.KV_CACHE_MULTIMODAL, max_seq_len=16384) - tokenizer = AutoTokenizer.from_pretrained(self.MODEL_PATH) - processor = AutoProcessor.from_pretrained( - self.MODEL_PATH, use_fast=True, trust_remote_code=True - ) - greedy = SamplingParams(max_tokens=96, temperature=0) - - def ask(text, items): - request = _media_request(tokenizer, processor, text, items) - return llm.generate_async(request, sampling_params=greedy).result().outputs[0].text - - with LLM( - self.MODEL_PATH, - speculative_config=self.MTP1, - **engine, - **_require_target(monkeypatch, **_VALIDATE), - ) as llm: - for name, (_, _, colour, directions) in _VIDEOS.items(): - answer = ask(_VIDEO_QUESTION, [("video", _synthetic_video(name))]).lower() - assert colour in answer and any(d in answer for d in directions), (name, answer) - - # Each label carries its own item's colour. - image = _synthetic_image((250, 200, 0)) - for items, video_colour in ( - ([("video", _synthetic_video("blue_leftward")), ("image", image)], "blue"), - ([("image", image), ("video", _synthetic_video("red_rightward"))], "red"), - ): - answer = ask(_MIXED_QUESTION, items).lower() - assert re.search(r"image:\W*yellow", answer), answer - assert re.search(rf"video:\W*{video_colour}", answer), answer - - with pytest.raises(RequestError, match="frames of 192x192 = 36864 patches"): - ask( - "What shape and colour is in the image?", - [("image", _synthetic_image(_RED, 3072))], - ) - answer = ask(_VIDEO_QUESTION, [("video", _synthetic_video("red_rightward"))]).lower() - assert "red" in answer and "right" in answer, answer diff --git a/tests/integration/defs/accuracy/references/acceptance_length.yaml b/tests/integration/defs/accuracy/references/acceptance_length.yaml index a87b9495bb69..37918c1460d9 100644 --- a/tests/integration/defs/accuracy/references/acceptance_length.yaml +++ b/tests/integration/defs/accuracy/references/acceptance_length.yaml @@ -135,7 +135,7 @@ TestGLM53FlashFP8::test_mtp: # every draft accepted. Each min_al below is a hand-set collapse tripwire unless its # comment says it is 95% of ref_al. -# test_mtp1_gsm8k_and_acceptance +# test_mtp1_gsm8k_mmmu_and_video, over gsm8k ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp1: # Smoke-scale, and labelled as such: measured on GB300 over the boot prompt # pair, STOCK 1.7403 against modeling_v2's 1.8063 on the same prompts. The @@ -154,18 +154,17 @@ ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp1: # measured, and tighten min_al then. min_al: 1.5 -# test_mtp3_cutedsl_bf16_ssm_state_gsm8k_and_acceptance: three drafts with the trunk's experts on -# CuteDSL over a bf16 recurrent state. -ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_cutedsl_bf16_state: - # modeling_v2's own full-gsm8k reading on GB300 (1319 samples); there is no built-in reading - # for this combination (the built-in model has no CuteDSL trunk). Each choice alone read - # 3.690 (CuteDSL) and 3.699 (bf16 state), against 3.704 for neither. +# test_mtp3_cutedsl_fp8_kv_bf16_state, over gsm8k: three drafts with the trunk's experts on +# CuteDSL over an fp8 KV cache and a bf16 recurrent state. +ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_cutedsl_fp8_kv_bf16_state: + # PENDING: modeling_v2's own full-gsm8k reading on GB300 (1319 samples); there is no built-in + # reading for this combination (the built-in model has no CuteDSL trunk). ref_al: 3.705 # 95% of ref_al, above the two-draft reading (~2.8): a third draft position that is never # accepted fails it. min_al: 3.520 -# test_mtp3_long_context_passkey[index_share]: three drafts with IndexShare on, bf16 KV, over +# test_mtp3_long_context_passkey: three drafts with IndexShare on, bf16 KV, over # 13 prompts of 3k-6k tokens (the sparse attention path). Most of the output is a copy of the # document's section headings, which the draft predicts almost always, so the readings sit # at the range's ceiling of 4.0. @@ -178,25 +177,20 @@ ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_long_context: # 3.760, a 6% drop the automatic floor would pass. min_al: 3.9 -# test_mtp3_long_context_passkey[fp8_kv]: the same prompts over an fp8 KV cache, IndexShare off. -ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_fp8_kv_long_context: - # 3.981 here; 3.973 with IndexShare on. +# test_mtp3_cutedsl_fp8_kv_bf16_state, over the pass keys of mtp3_long_context with IndexShare +# off. +ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_cutedsl_fp8_kv_bf16_state_long_context: + # PENDING ref_al: 3.981 # 2.5% below ref_al, as for mtp3_long_context. min_al: 3.88 -# test_mtp3_long_context_passkey[nvfp4_kv]: the same prompts over an NVFP4 KV cache, IndexShare off. -ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_nvfp4_kv_long_context: - # 3.971 here, against 4.000 over bf16 and 3.981 over fp8. - ref_al: 3.971 - # 2.5% below ref_al, as for mtp3_long_context. - min_al: 3.87 - -# test_mtp7_long_context_passkey: seven drafts with IndexShare on, over a bf16 recurrent state, -# over 12 prompts of 3k-12k tokens (the sparse attention path), whose output mostly copies the -# document's section headings: the readings sit near the range's ceiling of 8 (7.774 over an -# fp32 state). -ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp7_bf16_state_long_context: +# test_mtp7_long_context_passkey: seven drafts with IndexShare on, over an NVFP4 KV cache and a +# bf16 recurrent state, over 12 prompts of 3k-12k tokens (the sparse attention path), whose +# output mostly copies the document's section headings: the readings sit near the range's +# ceiling of 8 (7.857 over bf16 KV, 7.774 over bf16 KV and an fp32 state). +ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp7_nvfp4_kv_bf16_state_long_context: + # PENDING ref_al: 7.857 # 2.5% below ref_al, as for mtp3_long_context. min_al: 7.661 diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 480a51dc36cf..43e69d30002d 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -139,15 +139,14 @@ l0_b300: # stage instead), so 30 is 2.2x headroom. - unittest/_torch/modeling_v2 --ignore=unittest/_torch/modeling_v2/comm TIMEOUT (30) - accuracy/modeling_v2/test_gpt_oss.py::TestModelingV2GptOss120bSm10xDefault::test_gsm8k[tp1] - # qwen4_exp target, pre-merge: gsm8k; gsm8k with one MTP draft plus the - # acceptance-length floor, the only gate that sees a miscomputed draft layer; - # mmmu with MTP, which reaches the vision tower; and 3k-6k-token pass key prompts - # with three drafts and IndexShare over bf16 KV, the only prompts long enough to run - # the sparse attention path. The opt-in variants run post-merge (below). + # qwen4_exp target, pre-merge: gsm8k; one MTP engine for gsm8k with the acceptance-length + # floor (the only gate that sees a miscomputed draft layer), mmmu (the vision tower) and + # synthetic videos; and 3k-6k-token pass keys with three drafts and IndexShare over bf16 KV, + # the only prompts long enough to run the sparse attention path. The opt-in variants run + # post-merge (below). - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_gsm8k - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp1_gsm8k_and_acceptance - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp1_mmmu - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_long_context_passkey[index_share] + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp1_gsm8k_mmmu_and_video + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_long_context_passkey - condition: ranges: system_gpu_count: @@ -168,17 +167,11 @@ l0_b300: - unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k "TRTLLM" - unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k "CUTEDSL" - unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k "DEEPGEMM" - # qwen4_exp target, post-merge: the opt-in variants -- three drafts over CuteDSL experts and a - # bf16 recurrent state in one engine; fp8 and NVFP4 KV over the pass keys (their sparse - # attention path); seven drafts over 3k-12k-token pass keys on a bf16 state (the verify - # step's eight rows per request); seven drafts all rejected by forced acceptance (the - # speculative rollback end to end); synthetic videos; and the breakable prefill CUDA graph - # replayed against eager greedy continuations at the bucket edges and with context and - # generation rows in one batch. - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_cutedsl_bf16_ssm_state_gsm8k_and_acceptance - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_long_context_passkey[fp8_kv] - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_long_context_passkey[nvfp4_kv] + # qwen4_exp target, post-merge: the opt-in variants -- three drafts over CuteDSL experts, fp8 + # KV and a bf16 recurrent state in one engine (gsm8k and the pass keys); seven drafts over + # 3k-12k-token pass keys on NVFP4 KV and a bf16 state (the verify step's eight rows per + # request); and the breakable prefill CUDA graph replayed against eager greedy continuations + # at the bucket edges and with context and generation rows in one batch. + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_cutedsl_fp8_kv_bf16_state - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp7_long_context_passkey - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp7_forced_acceptance_extremes[all_rejected] - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp1_video - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_breakable_prefill_matches_eager From c9288d968c841342963e7479040fd0f07293a3b7 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Wed, 21 Oct 2026 02:57:23 -0700 Subject: [PATCH 26/28] [TRTLLM-17033][test] Qwen3.8-Flash-Next ops: at most three cells per op, one property and one guard test Each op certifies the shapes and layouts its forward reaches in at most three cells (four where the forward reaches four compiled variants: hc decode first, inject norm, qk norm rope, bf16 gemv, cached replay), keeps one property test and one single-case guard, and drops the wrong-variant, composition and determinism tests. Every op note lists what stays certified and what does not. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../modeling.py | 4 +- .../activation/fused_sigmoid_mul_inplace.py | 37 +- .../ops/activation/gelu_erf.py | 38 +- .../ops/activation/gelu_tanh.py | 35 +- .../ops/activation/silu_and_mul_rows.py | 8 +- .../ops/attention/qsa_canonical_selection.py | 27 +- .../ops/attention/qsa_expand_block_indices.py | 34 +- .../ops/attention/qsa_mtp_reuse_topk.py | 42 +- .../ops/attention/qsa_paged_index_scores.py | 83 +--- .../ops/attention/qsa_paged_index_store.py | 44 +- .../ops/attention/qsa_paged_kv_store.py | 36 +- .../ops/attention/qsa_paged_kv_store_nvfp4.py | 60 +-- .../ops/attention/qsa_paged_sparse_gqa.py | 68 ++- .../attention/qsa_paged_sparse_gqa_nvfp4.py | 91 +--- .../ops/attention/qsa_prefill_compress.py | 45 +- .../ops/attention/qsa_sparse_gqa_rowtile.py | 12 +- .../ops/attention/qsa_topk_expand_decode.py | 19 +- .../ops/attention/qsa_topk_select.py | 29 +- .../ops/attention/qwen4_exp_qk_norm_rope.py | 24 +- .../ops/attention/vision_rope_2d.py | 44 +- .../ops/embedding/ngram_hash.py | 3 +- .../ops/gemm/bf16_gemv_pdl.py | 18 +- .../ops/gemm/low_m_bf16_mm.py | 10 +- .../ops/hc/qwen4_exp_gated_residual.py | 36 +- .../ops/hc/qwen4_exp_hc_decode_first_cute.py | 26 +- .../hc/qwen4_exp_hc_decode_mix_fp4_cute.py | 32 +- .../ops/hc/qwen4_exp_hc_decode_tc_first.py | 10 +- .../ops/hc/qwen4_exp_hc_inject_norm.py | 37 +- .../ops/hc/qwen4_exp_hc_lowrank_act.py | 19 +- .../ops/hc/qwen4_exp_hc_up_gate_mean.py | 14 +- .../ops/hc/qwen4_exp_hc_up_gate_mean_fp4.py | 23 +- .../ops/hc/qwen4_exp_post_inject.py | 16 +- ...sl_nvfp4_gather_grouped_gemm_act_fusion.py | 50 +-- .../ops/moe/cute_dsl_nvfp4_grouped_gemm.py | 45 +- .../fp4_block_scale_moe_renorm_permuted.py | 8 +- .../moe/fp4_block_scale_moe_renorm_routed.py | 24 +- .../ops/moe/fp8_block_scale_moe_runner.py | 29 +- .../ops/moe/fused_sigmoid_gate_mul_add.py | 22 +- .../ops/moe/moe_sort.py | 38 +- .../ops/moe/moe_unpermute_inplace.py | 55 +-- .../ops/moe/renorm_moe_routing_op.py | 39 +- .../ops/norm/flashinfer_gemma_rmsnorm.py | 23 +- .../ops/norm/layer_norm.py | 22 +- .../ops/norm/rms_norm_gated_token_major.py | 27 +- .../ops/ple/gated_ngram_injection.py | 1 - .../ops/ple/ngram_token_history.py | 28 +- .../ops/ple/ple_gate_conv_inject_norm.py | 10 +- .../ops/ple/ple_ngram_embed.py | 23 +- .../ops/quantization/bf16_to_fp32.py | 1 - .../dequantize_e4m3_per_tensor.py | 23 +- .../ops/quantization/fp8_quantize_1x128.py | 43 +- .../ops/ssm/causal_conv1d_fwd.py | 54 +-- .../ops/ssm/causal_conv1d_triton_update.py | 101 ++--- .../ops/ssm/causal_conv1d_update.py | 52 ++- .../ops/ssm/chunk_gated_delta_rule.py | 66 +-- .../ops/ssm/cu_seqlens_from_seq_lens.py | 27 +- .../ops/ssm/dilated_short_conv_paged.py | 61 +-- .../ssm/extract_transpose_prefill_slice.py | 26 +- .../ssm/flashinfer_chunk_gated_delta_rule.py | 19 +- .../ops/ssm/fused_gdn_gating.py | 24 +- .../ops/ssm/fused_gdn_post_conv.py | 28 +- ...t_gated_delta_rule_cached_replay_update.py | 98 ++--- ...fused_recurrent_gated_delta_rule_update.py | 34 +- .../fused_sigmoid_gating_delta_rule_update.py | 50 +-- .../ops/ssm/gdn_decode_gated_norm_pdl.py | 17 +- .../ops/ssm/gdn_decode_rows_gated_norm.py | 23 +- .../ops/ssm/gdn_in_proj_conv1d_pdl.py | 11 +- .../ops/ssm/gdn_prefill_conv_post.py | 42 +- .../ops/ssm/pack_gdn_decode_qkv.py | 6 +- .../ops/ssm/reset_recurrent_state_rows.py | 22 +- .../ops/vision/patch_embed_conv3d.py | 6 +- .../ops/vision/pos_embed_interpolate.py | 37 +- ...t_modeling_v2_fused_sigmoid_mul_inplace.py | 71 +-- .../activation/test_modeling_v2_gelu_erf.py | 72 ---- .../activation/test_modeling_v2_gelu_tanh.py | 60 --- .../test_modeling_v2_silu_and_mul_rows.py | 13 +- ...est_modeling_v2_qsa_canonical_selection.py | 106 +---- ...st_modeling_v2_qsa_expand_block_indices.py | 175 +------- .../test_modeling_v2_qsa_mtp_reuse_topk.py | 272 ++---------- ...test_modeling_v2_qsa_paged_index_scores.py | 325 +------------- .../test_modeling_v2_qsa_paged_index_store.py | 144 +------ .../test_modeling_v2_qsa_paged_kv_store.py | 140 +----- ...st_modeling_v2_qsa_paged_kv_store_nvfp4.py | 142 +----- .../test_modeling_v2_qsa_paged_sparse_gqa.py | 388 +---------------- ...modeling_v2_qsa_paged_sparse_gqa_decode.py | 51 +-- ..._modeling_v2_qsa_paged_sparse_gqa_nvfp4.py | 123 +----- .../test_modeling_v2_qsa_prefill_compress.py | 203 +-------- ...test_modeling_v2_qsa_sparse_gqa_rowtile.py | 205 +-------- ...test_modeling_v2_qsa_topk_expand_decode.py | 119 ++---- .../test_modeling_v2_qsa_topk_select.py | 85 +--- ...test_modeling_v2_qwen4_exp_qk_norm_rope.py | 211 +-------- .../test_modeling_v2_vision_rope_2d.py | 212 +-------- .../embedding/test_modeling_v2_ngram_hash.py | 142 +----- .../gemm/test_modeling_v2_bf16_gemv_pdl.py | 37 +- .../gemm/test_modeling_v2_low_m_bf16_mm.py | 55 +-- ...st_modeling_v2_qwen4_exp_gated_residual.py | 101 +---- ...eling_v2_qwen4_exp_hc_decode_first_cute.py | 13 +- ...ing_v2_qwen4_exp_hc_decode_mix_fp4_cute.py | 27 +- ...odeling_v2_qwen4_exp_hc_decode_tc_first.py | 96 +---- ...st_modeling_v2_qwen4_exp_hc_inject_norm.py | 230 +--------- ...st_modeling_v2_qwen4_exp_hc_lowrank_act.py | 25 +- ...t_modeling_v2_qwen4_exp_hc_up_gate_mean.py | 69 +-- ...deling_v2_qwen4_exp_hc_up_gate_mean_fp4.py | 78 +--- .../test_modeling_v2_qwen4_exp_post_inject.py | 121 +----- ...sl_nvfp4_gather_grouped_gemm_act_fusion.py | 377 ++-------------- ...modeling_v2_cute_dsl_nvfp4_grouped_gemm.py | 403 ++---------------- .../test_modeling_v2_cute_dsl_weight_path.py | 30 +- ...ng_v2_fp4_block_scale_moe_renorm_routed.py | 29 +- ...est_modeling_v2_fp4_moe_renorm_cute_dsl.py | 36 +- ..._modeling_v2_fp8_block_scale_moe_runner.py | 356 +--------------- ..._modeling_v2_fused_sigmoid_gate_mul_add.py | 151 +------ .../moe/test_modeling_v2_moe_sort.py | 259 +---------- .../test_modeling_v2_moe_unpermute_inplace.py | 389 +++-------------- .../test_modeling_v2_renorm_moe_routing_op.py | 216 +--------- ...st_modeling_v2_flashinfer_gemma_rmsnorm.py | 69 +-- .../norm/test_modeling_v2_layer_norm.py | 187 -------- ..._modeling_v2_rms_norm_gated_token_major.py | 114 +---- .../test_modeling_v2_gated_ngram_injection.py | 131 +----- .../test_modeling_v2_ngram_token_history.py | 183 +------- ...t_modeling_v2_ple_gate_conv_inject_norm.py | 107 +---- .../ple/test_modeling_v2_ple_ngram_embed.py | 40 +- .../test_modeling_v2_bf16_to_fp32.py | 6 +- ..._modeling_v2_dequantize_e4m3_per_tensor.py | 113 +---- .../test_modeling_v2_fp8_quantize_1x128.py | 257 +---------- .../ssm/test_modeling_v2_causal_conv1d_fwd.py | 154 ++----- ...modeling_v2_causal_conv1d_triton_update.py | 166 ++------ .../test_modeling_v2_causal_conv1d_update.py | 137 ++---- ...test_modeling_v2_chunk_gated_delta_rule.py | 171 ++------ ...st_modeling_v2_cu_seqlens_from_seq_lens.py | 88 +--- ...st_modeling_v2_dilated_short_conv_paged.py | 280 +----------- ...ling_v2_extract_transpose_prefill_slice.py | 36 +- ...ng_v2_flashinfer_chunk_gated_delta_rule.py | 83 +--- .../ssm/test_modeling_v2_fused_gdn_gating.py | 140 +----- .../test_modeling_v2_fused_gdn_post_conv.py | 142 ++---- ...t_gated_delta_rule_cached_replay_update.py | 360 +--------------- ...fused_recurrent_gated_delta_rule_update.py | 394 ++--------------- ..._fused_sigmoid_gating_delta_rule_update.py | 245 ++--------- .../ssm/test_modeling_v2_gdn_block.py | 22 +- ...t_modeling_v2_gdn_decode_gated_norm_pdl.py | 67 +-- ..._modeling_v2_gdn_decode_rows_gated_norm.py | 88 +--- ...test_modeling_v2_gdn_in_proj_conv1d_pdl.py | 60 +-- .../test_modeling_v2_gdn_prefill_conv_post.py | 78 +--- .../test_modeling_v2_pack_gdn_decode_qkv.py | 60 +-- ..._modeling_v2_reset_recurrent_state_rows.py | 125 +----- .../test_modeling_v2_gdn_replay_verify.py | 12 +- .../test_modeling_v2_index_share.py | 2 +- .../test_modeling_v2_kv_range_probe.py | 12 +- .../test_modeling_v2_spec_rollback.py | 14 +- .../test_modeling_v2_patch_embed_conv3d.py | 120 +----- .../test_modeling_v2_pos_embed_interpolate.py | 131 +----- .../test_modeling_v2_vision_tower_calls.py | 23 +- .../test_modeling_v2_qwen4_exp_kv_quant.py | 53 +-- .../test_modeling_v2_qwen4_exp_moe_scalars.py | 30 +- ...modeling_v2_qwen4_exp_multimodal_embeds.py | 39 +- ...est_modeling_v2_qwen4_exp_pool_contract.py | 55 +-- ...eling_v2_qwen4_exp_thop_attention_calls.py | 59 ++- ..._modeling_v2_qwen4_exp_validate_request.py | 45 +- ...test_modeling_v2_qwen4_exp_weight_table.py | 22 +- 158 files changed, 1887 insertions(+), 11529 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py index 1da24685c9b5..14ace1194c5a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py @@ -3123,8 +3123,8 @@ def _gdn_verify( invisible to the acceptance rate. `pack_gdn_decode_qkv` does no Q/K L2 norm, hence `use_qk_l2norm_in_kernel=True` below.""" core = self.core - # [n*p, C] -> [n, C, p]: the multi-token conv reads channel-major (the - # transposed view is certified bitwise equal to the contiguous form). + # [n*p, C] -> [n, C, p]: the multi-token conv reads channel-major; its cells + # certify this transposed view directly. x3 = torch.transpose(torch.reshape(mq_d, [n, p, core.conv_dim]), 1, 2) conv_out = causal_conv1d_triton_update( x=x3, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/fused_sigmoid_mul_inplace.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/fused_sigmoid_mul_inplace.py index deee31706078..7987c692ddfd 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/fused_sigmoid_mul_inplace.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/fused_sigmoid_mul_inplace.py @@ -33,29 +33,17 @@ class FusedSigmoidMulInplace(OpWrapper): ARCHS = frozenset({Arch.SM_103}) CELLS: tuple[Cell, ...] = ( - Cell( - why="decode: 1 token, 24x256 gate as the strided half of q_proj (MTP)", - spec=dict(tokens=1, **_CKPT), - ), - Cell( - why="prefill 2048 tokens, strided q_proj half: >= 1024 rows picks 2048-wide blocks", - spec=dict(tokens=2048, **_CKPT), - ), - Cell( - why="decode 1 token: 2-D gate, the gate columns of the fused projection row", - spec=dict(tokens=1, heads=24, head_dim=256, dtype=torch.bfloat16, gate="fused_row"), - ), Cell( why="decode 128 tokens: 2-D gate off the fused projection row (stride 14080)", spec=dict(tokens=128, heads=24, head_dim=256, dtype=torch.bfloat16, gate="fused_row"), ), Cell( - why="prefill 8192 tokens: 2-D gate off the fused projection row", + why="prefill 8192 tokens: 2-D gate off the fused projection row, 2048-wide blocks", spec=dict(tokens=8192, heads=24, head_dim=256, dtype=torch.bfloat16, gate="fused_row"), ), Cell( - why="0 tokens: returns without launching", - spec=dict(tokens=0, **_CKPT), + why="MTP decode: 1 token, 24x256 gate as the strided half of q_proj", + spec=dict(tokens=1, **_CKPT), ), ) @@ -68,8 +56,8 @@ class FusedSigmoidMulInplace(OpWrapper): (an AssertionError with no message): `attention_output` 2-D with `stride(-1) == 1` (row stride free); `gate` 2-D of `attention_output`'s shape, or 3-D with `shape[0] == T` and `shape[1] * shape[2] == hidden`, - `stride(-1) == 1`. A CPU tensor raises ValueError from Triton. Hence no - `is_valid`. Cross-device operands were not probed. + `stride(-1) == 1`. A CPU tensor raises ValueError from Triton (measured). + Hence no `is_valid`. Cross-device operands were not probed. The pre-call contents of `attention_output` are destroyed; the gate is left bitwise unchanged. @@ -77,20 +65,21 @@ class FusedSigmoidMulInplace(OpWrapper): Numerics: bf16 bitwise equal to the fp32 reference except ~1e-6 of elements one ulp off (`tl.sigmoid`'s fast exp); default tolerance holds, nothing loosened. Feeding the q half instead of the gate half moves the - result by ~2.15 max abs at [17, 6144], so the per-head indexing is under - test. + result by ~2.15 max abs at [17, 6144] (measured), so the per-head indexing + is under test. Certified: bf16 operands at the checkpoint's hidden 6144, a compact - `attention_output`, and the target's two gate layouts (the strided q_proj - half and the 2-D window of the fused projection row). + `attention_output`; the 2-D window of the fused projection row at 128 + tokens (decode, 1024-wide blocks) and 8192 tokens (prefill, 2048-wide + blocks), and the strided q_proj half (MTP) at 1 token. Hard limits T * row stride token offset computed in int64 boundary 2^63 not probed (largest run 8192 rows) - Not certified: a strided `attention_output`, a gate dtype other than - `attention_output.dtype`, fp16/fp32 operands, other hidden sizes, - CUDA-graph capture, torch.compile. + Not certified: 0 tokens, the strided q_proj half above 1 token, a strided + `attention_output`, a gate dtype other than `attention_output.dtype`, + fp16/fp32 operands, other hidden sizes, CUDA-graph capture, torch.compile. """ def __call__(self, attention_output: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_erf.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_erf.py index edae62e6738d..07150279c931 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_erf.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_erf.py @@ -23,11 +23,6 @@ #: `merger.linear_fc1` is `[4608, 4608]`: spatial merge 2 over a 1152-wide tower. MERGER_WIDTH = 4608 -#: Rows the merger is asked for, `sum(t*h*w) / 4`: 1 the smallest, 140 a -#: 336x448 image (grid [1, 20, 28]) -- not a multiple of any MoE tile size -- -#: and 8192 the most one tower call carries (32768 patch rows). -MERGER_ROW_COUNTS = (1, 140, 8192) - _SQRT_2 = math.sqrt(2.0) @@ -46,20 +41,18 @@ class GeluErf(OpWrapper): # spec: shape, dtype; `dist` "uniform" over [lo, hi] (default) or "merger" # (a tower hidden state scaled by the merger norm's measured gamma, grouped # 4-to-1, through an fc1 at the checkpoint's measured weight/bias scale). - 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 - ) + ( + CELLS: tuple[Cell, ...] = ( Cell( why="activand built the way the merger builds it", spec=dict(shape=(512, MERGER_WIDTH), dtype=torch.bfloat16, dist="merger"), ), Cell( - why="float32 over [-3, 3], the erf/tanh divergence peak", - spec=dict(shape=(256, MERGER_WIDTH), dtype=torch.float32, lo=-3.0, hi=3.0), + why="merger activand, 8192 merged rows x 4608 bf16: the most one tower call carries", + spec=dict(shape=(8192, MERGER_WIDTH), dtype=torch.bfloat16, lo=-16.0, hi=16.0), + ), + Cell( + why="merger activand, 1 merged row x 4608 bf16", + spec=dict(shape=(1, MERGER_WIDTH), dtype=torch.bfloat16, lo=-16.0, hi=16.0), ), ) @@ -67,22 +60,23 @@ class GeluErf(OpWrapper): Metadata consumed: none. Out-of-place; `x` is not mutated. Certified on contiguous bf16 `[rows, 4608]` activands (what the merger's fc1 - returns) plus one fp32 cell. + returns) at 1, 512 and 8192 rows (8192, from 32768 patch rows, the most one + tower call carries). - Identity is pinned by fp32 only (the fp32 cell and the identity test). + Identity is pinned by fp32 only (the identity test). Max |erf GELU - tanh GELU| is 4.73e-4 at x = -2.699; in fp32 the tanh form lands ~1400x the correct path's deviation away, 31x outside the default gate. In bf16 one ulp at the output scale exceeds that gap (tanh variant 8.24e-3 vs correct 7.77e-3 max abs), so a bf16 match does not say which GELU was computed. - Values: `-inf -> nan` (the closed form does the same); `+inf -> +inf`; - large negatives give `-0.0`; `finfo.max` comes back unchanged. - CUDA-graph capturable; the output is a fresh allocation. + Measured, not certified: `-inf -> nan` (the closed form does the same); + `+inf -> +inf`; large negatives give `-0.0`; `finfo.max` comes back + unchanged; CUDA-graph capturable; the output is a fresh allocation. - Not certified: non-contiguous inputs, fp16, other widths, float8 - operands. No index-typed argument; numel past 2^31 - not probed. + Not certified: fp32 beyond the identity test, non-contiguous inputs, fp16, + other widths, float8 operands. No index-typed argument; + numel past 2^31 not probed. """ def __call__(self, x: torch.Tensor) -> torch.Tensor: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_tanh.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_tanh.py index 53db4ce79ebb..e9cd2ebec23a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_tanh.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_tanh.py @@ -41,24 +41,18 @@ class GeluTanh(OpWrapper): # spec: shape, dtype; `dist` "uniform" over [lo, hi] (default) or "fc1" (a # unit-normal hidden state through an fc1 at the checkpoint's measured # weight std 0.0197 / bias std 1.676). - CELLS: tuple[Cell, ...] = tuple( + CELLS: tuple[Cell, ...] = ( 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, 1024) - ) + ( + why="activand built through fc1 at the checkpoint's measured scale", + spec=dict(shape=(512, INTERMEDIATE_SIZE), dtype=torch.bfloat16, dist="fc1"), + ), Cell( why="block MLP activand at the tower call ceiling, 32768 patch rows x 4304 bf16", spec=dict(shape=(32768, INTERMEDIATE_SIZE), dtype=torch.bfloat16, lo=-16.0, hi=16.0), ), Cell( - why="activand built through fc1 at the checkpoint's measured scale", - spec=dict(shape=(512, INTERMEDIATE_SIZE), dtype=torch.bfloat16, dist="fc1"), - ), - Cell( - why="float32 over [-3, 3], the tanh/erf divergence peak", - spec=dict(shape=(256, INTERMEDIATE_SIZE), dtype=torch.float32, lo=-3.0, hi=3.0), + why="block MLP activand, 1 patch row x 4304 bf16", + spec=dict(shape=(1, INTERMEDIATE_SIZE), dtype=torch.bfloat16, lo=-16.0, hi=16.0), ), ) @@ -66,22 +60,23 @@ class GeluTanh(OpWrapper): Metadata consumed: none. Out-of-place; `x` is not mutated. Certified on contiguous bf16 `[rows, 4304]` activands (what the block MLP's - fc1 returns) plus one fp32 cell. + fc1 returns) at 1, 512 and 32768 rows (32768, the tower call ceiling). - Identity is pinned by fp32 only (the fp32 cell and the identity test). + Identity is pinned by fp32 only (the identity test). Max |tanh GELU - erf GELU| is 4.73e-4 (|x| ~ 2.7); in fp32 the erf form lands ~1600x the correct path's deviation away, 31x outside the default gate. In bf16 one ulp at the output scale exceeds that gap (erf variant 8.19e-3 vs correct 7.75e-3 max abs), so a bf16 match does not say which GELU was computed. - Values: `-inf -> nan` (the closed form does the same: -inf times a factor - saturated to 0); `+inf -> +inf`; large negatives give `-0.0`; `finfo.max` - comes back unchanged in bf16/fp16/fp32 (the cube is not taken in the input - dtype). CUDA-graph capturable; the output is a fresh allocation. + Measured, not certified: `-inf -> nan` (the closed form does the same: -inf + times a factor saturated to 0); `+inf -> +inf`; large negatives give + `-0.0`; `finfo.max` comes back unchanged in bf16/fp16/fp32 (the cube is not + taken in the input dtype); CUDA-graph capturable; the output is a fresh + allocation. - Not certified: non-contiguous inputs, fp16, other widths, float8 - operands. No index-typed argument; numel past 2^31 + Not certified: fp32 beyond the identity test, non-contiguous inputs, fp16, + other widths, float8 operands. No index-typed argument; numel past 2^31 not probed. """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/silu_and_mul_rows.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/silu_and_mul_rows.py index a79777be951b..cae68c64e1a9 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/silu_and_mul_rows.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/silu_and_mul_rows.py @@ -48,9 +48,8 @@ class SiluAndMulRows(OpWrapper): CELLS: tuple[Cell, ...] = ( Cell(why="prefill chunk 16384 rows of the fused 1288-wide output", spec=dict(rows=16384)), - Cell(why="decode batch 128, fused output", spec=dict(rows=128)), - Cell(why="1 row, fused output", spec=dict(rows=1)), - Cell(why="zero rows", spec=dict(rows=0)), + Cell(why="short prefill step, 128 rows of the fused output", spec=dict(rows=128)), + Cell(why="zero rows: no launch", spec=dict(rows=0)), ) note = """ @@ -66,6 +65,9 @@ class SiluAndMulRows(OpWrapper): Numerics: `g / (1 + exp(-g)) * u` in fp32 (Triton's `exp` and division), rounded once; bitwise `flashinfer_silu_and_mul` on the same values (16384 x 640 random bf16 values, the test's cells). + + Certified: bf16 at inter 640 on the 1288-wide fused output, 0, 128 and 16384 rows; + CUDA-graph replay at 128 rows. Not certified: fp16, other widths or row counts. """ def __call__(self, x: torch.Tensor, inter: int) -> torch.Tensor: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_canonical_selection.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_canonical_selection.py index ddd0b52d6281..b2736198b176 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_canonical_selection.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_canonical_selection.py @@ -114,18 +114,6 @@ class QsaCanonicalSelection(OpWrapper): ("chunk", 0), "relu", ), - _cell( - "windows 0, 1, 511, 512, 513, 1024, 2047, 2048, 2080", - 18, - ("ends", (0, 1, 511, 512, 513, 1024, 2047, 2048, 2080) * 2), - "relu", - ), - _cell( - "ties: a zero class of -0.0 and +0.0 that straddles k", - 16, - ("ends", (2048,) * 16), - "signed_zeros", - ), _cell( "input from the radix engine (qsa_topk_select), zero ties", 64, @@ -133,13 +121,6 @@ class QsaCanonicalSelection(OpWrapper): "relu_sparse", selection="radix", ), - _cell( - "CUDA graph: 32 rows captured, replayed with new scores and ends", - 32, - ("positions_random", 0, 8191), - "relu", - graph=True, - ), ) note = """ @@ -160,10 +141,10 @@ class QsaCanonicalSelection(OpWrapper): Cost: about 30 kernels -- elementwise passes, two int32 prefix sums over `[rows, cols]` and one sorted search over `[rows, k]` -- with no host read and no - data-dependent shape, so it captures into a CUDA graph and replays with new scores - and windows. It reads the whole `cols`-wide score row, not the live window: a caller - that knows a bound on the windows passes `logits[:, :bound]`, which gives the same - rows. A matrix of more than 2^24 elements is rewritten in row tiles. + data-dependent shape, so it can be captured into a CUDA graph (not certified: no + cell replays one). It reads the whole `cols`-wide score row, not the live window: a + caller that knows a bound on the windows passes `logits[:, :bound]`, which gives the + same rows. A matrix of more than 2^24 elements is rewritten in row tiles. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_expand_block_indices.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_expand_block_indices.py index bacb567ea3e0..ba1cf5acd32e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_expand_block_indices.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_expand_block_indices.py @@ -28,7 +28,6 @@ def _spec(batch, **overrides) -> dict: ("sweep", n) n rows at 4096 + 13 i, length = position + 1 ("serving", requests, context_lens) one packed step: context chunks, then two MTP rows per remaining request - ("rows", positions) explicit positions, length = position + 1 ("verify", starts, p) one MTP verify block of p rows per start, every row's length the block's end """ @@ -36,7 +35,6 @@ def _spec(batch, **overrides) -> dict: batch=batch, position_dtype=torch.int64, # qsa_logical_positions length_dtype=torch.int32, # qsa_sequence_lengths - layout="dense", **_CKPT, ) return {**base, **overrides} @@ -74,26 +72,18 @@ class QsaExpandBlockIndices(OpWrapper): CELLS: tuple[Cell, ...] = ( # ---- decode rows: the column-split launch below 128 rows, whole rows from 128 Cell(why="decode: one row, columns split across programs", spec=_spec(("sweep", 1))), - Cell(why="128 rows: the first row count on whole-row programs", spec=_spec(("sweep", 128))), # ---- MTP verify blocks (max_draft_len up to 7): p rows per request, every # row's length the block's end, so only the last row has length == position + 1. Cell( why="verify MTP7: 32 requests x 8 = 256 rows, whole-row programs", spec=_spec(("verify", _VERIFY_STARTS, 8)), ), - # ---- context rows: sequence_lengths > query_position + 1 inside chunks + # ---- context rows: sequence_lengths > query_position + 1 inside chunks; chunk + # 0 starts at position 0, so its rows also carry every tail width 0..3. Cell( why="warm-up capture: 2 x 4096 rows, the [8192, 512] buffer's height", spec=_spec(("serving", 2, (4096, 4096))), ), - Cell( - why="tail widths 0..3 across 4096, and positions 0, 3, 7", - spec=_spec(("rows", (4095, 4096, 4097, 4098, 4099, 7, 3, 0))), - ), - Cell( - why="a row slice of the engine's taller qsa_topk_indices buffer", - spec=_spec(("sweep", 4), layout="row_slice"), - ), ) note = """ @@ -103,8 +93,8 @@ class QsaExpandBlockIndices(OpWrapper): request's whole KV length broadcast over its rows, so inside a context chunk it exceeds `position + 1` on every row but the last). - Preconditions that are values, so not `is_valid` guards (both measured to - fail silently): + Preconditions that are values, so not `is_valid` guards (both fail + silently; no test pins them): - `block_indices` is prefix-packed: every group id before every `-1`. The tail's column comes from the *count* of valid groups, so `[-1, 3]` at position 15, length 16, cr 4, topk 8 returns all `-1` -- tokens @@ -115,16 +105,16 @@ class QsaExpandBlockIndices(OpWrapper): inside `sequence_lengths` is expanded verbatim. Duplicate group ids expand to duplicate tokens. - Every metadata violation is rejected loudly by the callable (rank, width - vs geometry, int32-only `block_indices`, int32/int64 metadata, one entry - per row, `stride(1) == 1`, contiguous metadata, CUDA), so `is_valid` adds - nothing. A one-element strided metadata view reports contiguous and is - accepted -- benign at one row. + The callable rejects metadata violations loudly (rank, width vs geometry, + int32-only `block_indices`, int32/int64 metadata, one entry per row, + `stride(1) == 1`, contiguous metadata, CUDA), so `is_valid` adds nothing; + the test pins the width-vs-geometry rejection. A one-element strided + metadata view reports contiguous and is accepted -- benign at one row. Launch: one program per (row, 256 columns) below 128 rows, one per row - from 128 up; 127 and 128 rows agree on their shared rows. No autotune. - CUDA-graph capturable (captured at 4 and at 256 rows); a replay recomputes - from the live buffers, and the captured row count fixes the launch shape. + from 128 up. No autotune. CUDA-graph capturable (captured at 256 rows); a + replay recomputes from the live buffers, and the captured row count fixes + the launch shape. Hard limits (read from the kernel, not probed): the row offset is `program_id * stride(0)` in int32, so `rows * stride(0) >= 2^31` (4 Mi diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_mtp_reuse_topk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_mtp_reuse_topk.py index 89367af18c42..582e514682ea 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_mtp_reuse_topk.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_mtp_reuse_topk.py @@ -16,10 +16,8 @@ #: `indexer_compress_ratio` 4. _BLOCK_TOPK = 512 #: `(slots, this layer's slot)` of the shared buffers: with the MTP layer in the -#: target's cache manager, 12 trunk full-attention layers plus the MTP layer; a -#: separate draft manager owns one slot, the MTP layer's. +#: target's cache manager, 12 trunk full-attention layers plus the MTP layer. _COMBINED_SLOTS = (13, 12) -_DRAFT_SLOTS = (1, 0) def _spec(contexts, **overrides) -> dict: @@ -28,11 +26,8 @@ def _spec(contexts, **overrides) -> dict: `contexts` says what each request's captured row and current row are: ("mixed", B) B requests; the first 16 at fixed boundary positions (empty, short, 511, 512 and full rows), the rest seeded - ("short", B) captured positions <= 2046: fewer than 512 valid groups ("full", B) captured positions >= 2047: all 512 slots valid ("explicit", P) one request per captured position in P - ("gaps", pairs) (captured_visible, current_visible) per request, set - directly rather than through positions A request captured at position `p` (its last accepted row) holds `(p + 1) // 4` complete groups; its draft row `d` positions later sees `(p + d + 1) // 4`. `drafted` cycles over the requests. @@ -41,11 +36,9 @@ def _spec(contexts, **overrides) -> dict: contexts=contexts, drafted=(1, 2, 3), rows_per_request=1, - request_map="identity", shared_rows=64, slots=_COMBINED_SLOTS, block_topk=_BLOCK_TOPK, - layout="engine", metadata_dtype=torch.int32, graph_batches=(), ) @@ -90,15 +83,13 @@ class QsaMtpReuseTopk(OpWrapper): ARCHS = frozenset({Arch.SM_103}) CELLS: tuple[Cell, ...] = ( - # ---- the two behaviours, at every boundary ------------------------- + # ---- decode batch: one row per request; the first 16 requests sit at the + # boundaries (empty, short, 511-id, 512-id and full rows), fresh 0 and 1. Cell( - why="8 requests, one per boundary: 0/1/100/511/512/2047 groups, fresh 0 and 1", - spec=_spec( - ("explicit", (2, 3, 402, 2046, 2050, 8188, 2051, 2047)), - drafted=(1, 1, 1, 1, 1, 3, 3, 1), - shared_rows=8, - ), + why="decode MTP: 32 requests (envelope), 256-row shared buffer", + spec=_spec(("mixed", 32), shared_rows=256), ), + # ---- the riskiest edge: two fresh groups displace a full row's two oldest Cell( why="8 requests at fresh 2: full, 511-id and short rows gain two groups", spec=_spec( @@ -107,15 +98,6 @@ class QsaMtpReuseTopk(OpWrapper): shared_rows=8, ), ), - # ---- decode batches: one row per request ---------------------------- - Cell( - why="decode MTP: 32 requests (envelope), 256-row shared buffer", - spec=_spec(("mixed", 32), shared_rows=256), - ), - Cell( - why="permuted request map, draft manager's single slot", - spec=_spec(("mixed", 32), request_map="permuted", slots=_DRAFT_SLOTS, shared_rows=32), - ), # ---- CUDA graph ----------------------------------------------------- Cell( why="CUDA graph, drafted 1..6: captured at 1, 8 and 32 requests, new-data replay", @@ -154,7 +136,7 @@ class QsaMtpReuseTopk(OpWrapper): - `shared_visible_blocks[q]` is the visible count of the very row that was captured, and `visible_blocks[r]` is `(position + 1) // cr` of row `r`. - Loud rejections: `ValueError` when the shared width differs from the output + Loud rejections (not pinned by a test): `ValueError` when the shared width differs from the output width, when `request_indices` or `visible_blocks` does not hold one element per output row, or on CPU tensors; a 1-D `output` or `shared_indices`; a width that is not a power of two (Triton compilation). @@ -163,11 +145,13 @@ class QsaMtpReuseTopk(OpWrapper): (the kernel addresses raw pointers densely and ignores strides); a 3-D `shared_indices` whose middle extent equals `block_topk` (read as slot 0); a 2-D `shared_visible_blocks` (the whole `[slots, seqs]` buffer is read as - slot 0's counts). + slot 0's counts). The test pins the last. - Certified: int32 ids and metadata, `block_topk` 512, up to 32 requests. - No host synchronization and shape-only validation, so it captures into a - CUDA graph. + Certified: int32 ids and metadata, `block_topk` 512, up to 32 requests, the + identity request map over the target cache manager's 13-slot buffers, CUDA + graph capture at 1, 8 and 32 requests with a new-data replay. + Not certified: a permuted request map, a separate draft manager's one-slot + buffers. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_scores.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_scores.py index fc41479f5980..fff162545bc7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_scores.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_scores.py @@ -70,7 +70,7 @@ class QsaPagedIndexScores(OpWrapper): that divides 128 into more than one row (1, 2, 4, ..., 64) it selects a query-tiled kernel sharing one key gather across `128 // heads` rows; otherwise the row-wise kernel runs. Each cell names the kernel it was measured on. The - two agree to one fp32 ulp, not bitwise. + two are not bitwise equal. Fusion boundary: the table walk, the gather, the dots, relu, head sum, scale and causal bound are inside. Producing `q` (norm, RoPE), writing and @@ -95,32 +95,6 @@ class QsaPagedIndexScores(OpWrapper): **_TRUNK_POOL_NVFP4, ), ), - Cell( - why="NVFP4 trunk plane, page stride 49344: prefill 300 rows from 0", - spec=_spec( - layout="capture", - lengths=(300,), - slots=8, - pages=_WIDTH, - only_visible_blocks=True, - context_rows=True, - branch="query_tile", - **_TRUNK_POOL_NVFP4, - ), - ), - 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 MTP7: 32 slots x 8 rows = 256 rows, the last decode-sized count", spec=_spec( @@ -146,32 +120,6 @@ class QsaPagedIndexScores(OpWrapper): branch="query_tile", ), ), - Cell( - why="mixed step: 195-token chunk + 31 MTP decodes = 257 rows: query-tiled", - 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 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", - ), - ), ) note = """ @@ -188,22 +136,22 @@ class QsaPagedIndexScores(OpWrapper): run: an uncompressed anchor scores as a raw index key. `rows == 0` returns an empty `[0, columns]` tensor without launching. - The op validates its own metadata loudly (`ValueError`s; one per refusal - family is pinned by the test), so `is_valid` adds nothing. + The op validates its own metadata loudly (`ValueError`s; the test pins the + one-entry-per-row metadata refusal), so `is_valid` adds nothing. - `only_visible_blocks=True` really leaves garbage: with a poisoned - allocation freed just before the call, the whole unscored tail came back - poisoned. At 8191 rows from position 0 over 2080 columns that is a full + `only_visible_blocks=True` really leaves garbage (measured, not pinned by a + test): with a poisoned allocation freed just before the call, the whole + unscored tail came back poisoned. At 8191 rows from position 0 over 2080 columns that is a full 32-column band (2048..2079) for every row, 8.65M cells. A consumer that scans `[row_start, row_end)` -- `qsa_topk_select` -- is unaffected. fp32 operands are accepted but `tl.dot` rounds them to a tensor-core - format: 8.1e-3 max abs against an fp32 reference on scores to 7.8, outside - the test's band. No cell covers fp32; the side cache is bf16. + format: 8.1e-3 max abs against an fp32 reference on scores to 7.8 (measured), + outside the test's band. No cell covers fp32; the side cache is bf16. - CUDA graphs: 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. + CUDA graphs: the test captures the decode launch (`use_pdl`, `prefetch`) at + 4 rows and replays it after `q` and the anchor keys were rewritten in + place; every replay matches an eager call bitwise. `prefetch` (with `use_pdl` and `only_visible_blocks`, at generation-sized batches) reads the pages and compressed keys before the PDL wait and q after @@ -214,13 +162,14 @@ class QsaPagedIndexScores(OpWrapper): Every cell runs `only_visible_blocks=True` (the target never passes False), 4 index heads x 128, compress ratio 4, 32-token pages, bf16 operands, int64 - positions and int32 request indices. + positions and int32 request indices. The NVFP4-pool page stride is + certified on the row-wise kernel only. Not certified: more than 8191 rows, tables wider than 260 pages, more than 32 table rows, any other head geometry, compress ratio, page size or - operand dtype, and the `-inf` padded mode as a cell (the test gates it - against the reference outside the cells). Hard limits: not probed; the largest - certified extent is 8191 x 2080 output elements. + operand dtype, a context step's mixed request rows, and the `-inf` padded + mode. Hard limits: not probed; the largest certified extent is 8191 x 2080 + output elements. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_store.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_store.py index 6bea19a8b576..0100f49724d9 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_store.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_store.py @@ -135,13 +135,10 @@ def _launch_rows(rows: int) -> Tuple[int, int]: hole=None, ) -_DECODE = ((0, 1), (1, 1), (31, 1), (32, 1), (33, 1), (63, 1), (64, 1), (95, 1)) - #: The target's serving envelope: ISL 8192 + OSL 1024 puts decode rows up to position #: 9215 (page column 287), 128 concurrent requests, and `max_num_tokens` 16384 #: rows in one prefill or mixed step. _WIDE = dict(slots=160, width=300) -_DECODE_128 = tuple(((r * 72 + 9) % 9216, 1) for r in range(128)) _PREFILL_16K = ((0, 8192), (0, 8192)) #: 32 requests of one MTP7 verify block each (8 rows), at starts that put some blocks across a @@ -192,11 +189,6 @@ class QsaPagedIndexStore(OpWrapper): "launch, blocks straddling pages", spec=_spec(_MTP7_32, pages=1200, slots=40), ), - Cell( - why="decode M=128 at positions up to 9153 (page column 286), token_k off the " - "fused projection row", - spec=_spec(_DECODE_128, pages=18432, key_layout="fused_row", **_WIDE), - ), Cell( why="prefill 16384 rows (max_num_tokens), 16 rows per program, token_k off the " "fused projection row", @@ -204,7 +196,8 @@ class QsaPagedIndexStore(OpWrapper): ), Cell( why="NVFP4-KV side planes: index and position planes as strided views of one " - "per-slot buffer holding 12 index-K pages and the position page", + "per-slot buffer holding 12 index-K pages and the position page; one row " + "on an unallocated (-1) table entry is skipped", spec=_spec( ((0, 1), (31, 1), (32, 1), (33, 40)), pages=8, @@ -212,16 +205,9 @@ class QsaPagedIndexStore(OpWrapper): layers=12, page_stride=49344, position_strides=(24672, 3, 1), + hole="unallocated_page", ), ), - Cell( - why="skip: table entry -1 (unallocated)", - spec=_spec(_DECODE, pages=32, hole="unallocated_page"), - ), - Cell( - why="skip: negative request index", - spec=_spec(_DECODE, pages=32, hole="negative_request"), - ), ) note = """ @@ -246,7 +232,8 @@ class QsaPagedIndexStore(OpWrapper): - destination slots are unique across the batch: two rows naming one (page, slot) both store, unordered. The engine's mapping cannot collide. - Loud rejections (`ValueError`, before launching): metadata that is not + Loud rejections (`ValueError`, before launching; the test pins the + unequal-length metadata one): metadata that is not one-dimensional, of unequal length, or not int32/int64; `token_k` or the coordinates with a row count other than the metadata's; an index plane whose page size, head count or head width disagrees with @@ -254,9 +241,9 @@ class QsaPagedIndexStore(OpWrapper): disagrees or that is narrower than the coordinates; a table that is not a 2-D integer tensor. `rows == 0` returns before launching. - Values the caller owns -- data, not metadata: coordinates wider than the - plane's dtype wrap (int64 `2^31 + 5` stores as -2147483643); floats - truncate toward zero. + Values the caller owns -- data, not metadata (measured, not pinned by a + test): coordinates wider than the plane's dtype wrap (int64 `2^31 + 5` + stores as -2147483643); floats truncate toward zero. One launch per call; HEAD_DIM, the head count, the page size and the coordinate width are constexprs, everything else (strides, pool and table @@ -265,19 +252,20 @@ class QsaPagedIndexStore(OpWrapper): (strides still specialize on 1 and 16-divisibility). Up to 256 rows run one row per program (1 warp), above that 16 rows per program (4 warps). CUDA-graph capturable (captured at 256 rows, 32 MTP7 verify blocks); a - replay follows positions rewritten in place. Repeating a call is - idempotent. + replay follows positions rewritten in place. Hard limits: every index is widened to int64 inside the kernel before any arithmetic; an int64 position at `2^32 + 5` lands past the table and - is skipped rather than truncated. Rows are a grid axis: 2^31 programs. + is skipped rather than truncated (read from the kernel, not pinned by a + test). Rows are a grid axis: 2^31 programs. - Skips certified by a cell: an unallocated (`-1`) table entry and a negative - request index; the test also drives an int64 position past the table. The - other skip conditions share that mask and have no cell. + Skip certified by a cell: an unallocated (`-1`) table entry, one row of the + NVFP4-plane cell. The other skip conditions (a negative request index among + them) share that mask and have no cell. Not certified: an fp8 index plane, keys other than bf16, any index geometry - other than 1 head x 128 at 32-token pages, strided metadata or tables, + other than 1 head x 128 at 32-token pages, decode batches other than 32 MTP7 + verify blocks, strided metadata or tables, stride-0 coordinates, multiple devices, concurrent streams onto one pool, duplicate destinations. """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store.py index 9fb56a0e03c8..be0791c3ccd6 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store.py @@ -17,9 +17,7 @@ #: 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. `values` -#: fills the rows: `every_bf16` is all 65536 bf16 bit patterns (needs 128 rows -#: of 2 x 256). +#: max_blocks_per_seq). `pages` is the pool, independent of the batch. _CKPT = dict( hkv=2, head_dim=256, @@ -31,13 +29,10 @@ cache_dtype=torch.bfloat16, request_dtype=torch.int32, position_dtype=torch.int64, - k_layout="qkv_slice", - table_layout="dense", hole=None, # Heads of further columns after v in each row: 27 is the target's fused input # projection row [q | k | v | gate | index | pad] (55 heads of 256). row_pad_heads=0, - values="randn", ) _F8 = torch.float8_e4m3fn @@ -91,30 +86,20 @@ class QsaPagedKvStore(OpWrapper): # targets split them off a packed QKV row. Cell( why="decode: 8 rows off the fused projection row (row stride 14080), " - "at aligned, mid and page-crossing positions", - spec=_spec(_DECODE, pages=64, row_pad_heads=27), + "at aligned, mid and page-crossing positions, plus one row on an " + "unallocated (-1) table entry that is skipped", + spec=_spec(_DECODE, pages=64, row_pad_heads=27, hole="unallocated_page"), ), Cell( why="prefill 1024 rows off the fused projection row (row stride 14080)", spec=_spec(((0, 256),) * 4, pages=160, row_pad_heads=27), ), # MTP7 verify blocks: 8 rows per request, so one block spans up to two pages. - Cell( - why="decode MTP7: 32 requests x 8 rows (256), blocks straddling pages", - spec=_spec(_MTP7_32, pages=1200, slots=40), - ), + # The non-default pool the forward reaches: fp8-e4m3. Cell( why="fp8 decode MTP7: 32 requests x 8 rows (256), blocks straddling pages", spec=_spec(_MTP7_32, pages=1200, slots=40, cache_dtype=_F8), ), - Cell( - why="fp8 conversion: all 65536 bf16 bit patterns, as torch casts them", - spec=_spec(((0, 128),), pages=8, cache_dtype=_F8, values="every_bf16"), - ), - Cell( - why="skip: table entry -1 (unallocated)", - spec=_spec(_DECODE, pages=64, hole="unallocated_page"), - ), ) note = """ @@ -138,7 +123,7 @@ class QsaPagedKvStore(OpWrapper): `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. + table (not pinned by a test). `rows == 0` returns before launching. Strides and the pool/table extents are Triton constexprs: each distinct layout or pool size compiles its own kernel. CUDA-graph capturable @@ -146,12 +131,13 @@ class QsaPagedKvStore(OpWrapper): rewritten in place. Hard limits: both metadata values are cast to int32 inside the kernel - (`.to(tl.int32)`); an int64 position at `2^32 + 5` was probed and lands - on slot 5 of page column 0 -- written, not skipped. The engine's - positions stay below 2^14. + (`.to(tl.int32)`); an int64 position at `2^32 + 5` was probed once (not + pinned by a test) and lands on slot 5 of page column 0 -- written, not + skipped. The engine's positions stay below 2^14. Not certified: duplicate destinations, multiple devices, a scaled fp8 - pool (there is no scale argument), fp16 rows, head geometry other than + pool (there is no scale argument), fp8 conversion of values past 448, the + packed-QKV-row (MTP) layout over a bf16 pool, fp16 rows, head geometry other than this checkpoint's (D 256, Hkv 2, 32-token pages), int32 positions or int64 request indices, head-major rows, sliced or transposed tables, `rows == 0`, and of the rows that resolve to no slot any but an diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store_nvfp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store_nvfp4.py index 12a593a90d41..73e6ec1bebf7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store_nvfp4.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store_nvfp4.py @@ -53,7 +53,6 @@ global_v=1.0, values="randn", hole=None, - graph_batches=(), thop_append=False, slot_pages=2, layer=0, @@ -377,51 +376,21 @@ class QsaPagedKvStoreNvfp4(OpWrapper): CELLS: tuple[Cell, ...] = ( Cell( - why="decode: 8 rows at aligned, mid-page and page-crossing positions", - spec=_spec(_DECODE, pages=64), - ), - Cell( - why="prefill: 1024 rows, 4 requests x 256 walking 8 pages each", - spec=_spec(((0, 256),) * 4, pages=160), - ), - Cell( - why="every finite bf16 value, K in bit-pattern order and V shuffled across blocks", - spec=_spec(((0, 128),), pages=8, values="finite_bf16"), + why="decode: 8 rows at aligned, mid-page and page-crossing positions, with zeros, " + "-0.0 and a 900.0 outlier inside blocks, and a row on an unallocated (-1) page: " + "zero scales, saturation and the skip", + spec=_spec(_DECODE, pages=69, slots=8, values="edge", hole="unallocated_page"), ), Cell( - why="zeros, -0.0 and a 900.0 outlier inside blocks, and a row on an unallocated " - "(-1) page: zero scales, saturation and the skip", - spec=_spec(_DECODE, pages=69, slots=8, values="edge", hole="unallocated_page"), + why="prefill: 1024 rows, 4 requests x 256 walking 8 pages each, in the engine's " + "trunk layout: layer 11's planes among the 24 role pages of a slot", + spec=_spec(((0, 256),) * 4, pages=160, slot_pages=24, layer=11), ), Cell( why="the attention op's own append over the same rows (context, chunk, decode " "P=1 and P=3) at unit second-level scales: data and both scale layouts bitwise equal", spec=_spec(((0, 115), (0, 101), (0, 36)), pages=96, thop_append=True), ), - Cell( - why="the engine's trunk layout: layer 11's planes among the 24 role pages of a slot", - spec=_spec(((0, 64), (1000, 2), (31, 2)), pages=96, slots=8, slot_pages=24, layer=11), - ), - # Triton compiles once per class of each plain-int argument -- equal to 1, a multiple - # of 16, or anything else. The pool's slot count and the table's height are the two a - # served call varies; with the cells above (pool, height) is (16k, 16k), (other, 16k), - # (other, other) and (16k, other), and these two add table height 1. - Cell( - why="decode batch 1: slot count other, table height 1", - spec=_spec(((0, 1),), pages=69, slots=1), - ), - Cell( - why="decode batch 1: slot count a multiple of 16, table height 1", - spec=_spec(((0, 1),), pages=64, slots=1), - ), - Cell( - why="decode MTP graph: G = 1, 3 and 8 requests x 8 rows, replayed bitwise equal to eager", - spec=_spec( - tuple((29 + 61 * i, 8) for i in range(8)), - pages=256, - graph_batches=(1, 3, 8), - ), - ), ) note = """ @@ -441,7 +410,8 @@ class QsaPagedKvStoreNvfp4(OpWrapper): K scales linear, V scales in four-token groups. The attention op's append, run over the same rows, leaves bitwise the same pool (a cell), at the unit second-level scales the target binds. - The callable rejects loudly (ValueError): non-matching or non-3-D k/v, a row + The callable rejects loudly (ValueError; the test pins the non-contiguous + metadata refusal): non-matching or non-3-D k/v, a row dtype outside bf16/fp16, a head_dim that is not a power of two >= 32, a tokens_per_block that is not a multiple of 4, caches that are not uint8 of the stated shapes or whose last dim or scale head page is strided, a non-2-D @@ -451,7 +421,8 @@ class QsaPagedKvStoreNvfp4(OpWrapper): Every offset is formed in 64 bits: rows, heads, requests, positions and pages are widened before any product. CUDA-graph capturable; a replay follows rows, - pages and positions rewritten in place. + pages and positions rewritten in place (the test replays 8 MTP verify blocks of + 8 rows bitwise equal to eager). Compiled variants: the page count (`num_pages`, the pool's slot count), the table's height (`num_requests`) and width and every stride are plain-int @@ -460,8 +431,9 @@ class QsaPagedKvStoreNvfp4(OpWrapper): tokens_per_block. A served call puts the width and each stride in one fixed class (the k/v head slices of the packed QKV row, the planes' slot, head and token strides, the contiguous table's) and varies the page count and the - table height; the cells run both reachable page-count classes against every - table-height class. + table height. The cells run (other, other) and (multiple of 16, multiple of + 16); the other combinations, a table height of 1 among them, are not + certified. Not finite values: a NaN or +-inf in a block is not certified (the e2m1 and e4m3 conversions saturate, the block maximum is the kernel's `max`). @@ -518,8 +490,8 @@ def reference( Quantization in fp32 torch with IEEE reciprocals where the kernel uses `rcp.approx.ftz`; the two agree on every operand the cells drive (6.0, the - unit second-level scale, and every e4m3 scale value), which the cells' bitwise - gate is what pins. + unit second-level scale, and the e4m3 scale values their rows produce), which + the cells' bitwise gate is what pins. """ out = [ t.clone(memory_format=torch.contiguous_format) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa.py index da6e70489838..21609fd37846 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa.py @@ -28,8 +28,6 @@ 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; # "fused_row" is the target's fused input projection row, 55 heads of 256. q_layout="qkv_slice", @@ -58,11 +56,6 @@ def _decode(segments, launches=_SPLIT_K, **overrides) -> dict: return _spec(**{**base, **overrides}) -def _mixed(requests, context_lens, **overrides) -> dict: - base = dict(mixed=(requests, context_lens), slots=32, prefill_pair=True, launches=_PREFILL) - return _spec(**{**base, **overrides}) - - class QsaPagedSparseGqa(OpWrapper): """Sparse GQA over the selected tokens of an HND paged K/V pool. @@ -116,30 +109,15 @@ class QsaPagedSparseGqa(OpWrapper): CELLS: tuple[Cell, ...] = ( # ---- generation: no prefill pair, split-K --------------------------- Cell( - why="decode MTP7: 32 slots x 8 rows = 256 rows, 512 programs, 4-way split-K", - spec=_decode(tuple((100 + (499 * g) % 8000, 8) for g in range(32))), - ), - Cell( - why="decode: 1 row, gate off the fused row, in the merge, PDL", - spec=_decode(((4000, 1),), q_layout="fused_row", gate="fused_row"), - ), - Cell( - why="decode: 128 rows, 8-way split-K, gate fused into the merge", - spec=_decode(((600, 1),) * 128, q_layout="fused_row", slots=160, gate="fused_row"), - ), - Cell( - why="300 rows, no prefill pair: split-K, one split, no merge", - spec=_decode(((600, 1),) * 300, launches=("splitk",), slots=384), - ), - Cell( - why="holes: all-padding row, foreign request, unmapped page", - spec=_decode(((900, 1),) * 4, holes=True), + why="decode MTP7: 32 slots x 8 rows = 256 rows, 512 programs, 4-way split-K, " + "q and the gate off the fused row, the gate fused into the merge", + spec=_decode( + tuple((100 + (499 * g) % 8000, 8) for g in range(32)), + q_layout="fused_row", + gate="fused_row", + ), ), # ---- context: the prefill pair is passed ---------------------------- - Cell( - why="fp8 context 128 rows from 2032: pair passed, still split-K", - spec=_decode(((2032, 128),), prefill_pair=True, **_FP8), - ), Cell( why="fp8 context 272 rows from 2032: causal-prefix prefill kernel", spec=_decode(((2032, 272),), launches=_PREFILL, prefill_pair=True, **_FP8), @@ -168,29 +146,35 @@ class docstring). QSA's producer -- `qsa_topk_select` with zero is not an `is_valid` guard. - The prefill pair is a property of the launch: passed because some rows are context, it puts the generation rows on the prefill kernel - too, and the bound applies to them (asserted with a planted index). + too, and the bound applies to them (measured with a planted index; not + pinned by a test). - There is no KV scale argument: an fp8 pool is read at face value (each code widens exactly to bf16 under a bf16 q), so it must hold K/V written at unit scale. - `D` must be a power of two (192 is refused); `Hq % Hkv == 0`. - The op validates its own inputs loudly (`ValueError`s, one per guard - family pinned by the test, plus a selection width the compress ratio - cannot produce, refused on the prefill kernel only), so `is_valid` adds - nothing. `rows == 0` returns an empty `[0, Hq, D]` tensor. + The op validates its own inputs loudly (`ValueError`s; the test pins the + refusal of positions passed without a compress ratio; a selection width + the compress ratio cannot produce is refused on the prefill kernel only), + so `is_valid` adds nothing. `rows == 0` returns an empty `[0, Hq, D]` + tensor (not certified). PDL: `TRTLLM_ENABLE_PDL` (default on) enables programmatic dependent - launch for multi-split launches of 4 rows or fewer; both launch modes are - inside the cells. Calls are bitwise deterministic. + launch for multi-split launches of 4 rows or fewer; no cell runs one (the + target sends gated decode batches that small to + `qsa_paged_sparse_gqa_decode`). Calls are bitwise deterministic (pinned + on a 32-split launch and its merge). The decode site carries up to 8 rows per request (max_draft_len 7): the - cells run 32 requests of 8 rows, 512 `rows * Hkv` programs and a 4-way - split-K launch; the 1- and 128-row decode cells run more splits. + decode cell runs 32 requests of 8 rows, 512 `rows * Hkv` programs and a + 4-way split-K launch with the output gate in its merge. Not certified: capture in a CUDA graph (the split-K path allocates its - partials inside the call); head geometries other than 24/2 x 256; - `tokens_per_block` other than 32; more than 8191 rows; the split-K kernel - above 300 rows; an fp16 pool or q; standalone contiguous K/V buffers; - int64 request indices. + partials inside the call); PDL launches; head geometries other than 24/2 + x 256; `tokens_per_block` other than 32; more than 8191 rows; the split-K + kernel above 256 rows, with one split (no merge) or without the gate; an + fp8 pool on the split-K kernel; rows whose selection resolves no token + (all padding, a foreign request, unmapped pages); an fp16 pool or q; + standalone contiguous K/V buffers; int64 request indices. Hard limits: not probed; the largest certified extent is 8191 rows x 2051 selected tokens through a [32, 260] table. """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4.py index 1bfb5fd146b4..6909d8cde45d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4.py @@ -57,8 +57,7 @@ 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, or `mixed` is - `(requests, context chunk lengths)`; `prefill_pair` says whether `query_positions` + `segments` is `(first position, rows)` per request, in row order; `prefill_pair` says whether `query_positions` / `compress_ratio` are passed, which the target does exactly when the step has a context request; `pool_pages` fixes the pool's page count (the requests then sit on its last pages); `capture` sweeps the decode site under a CUDA graph. @@ -71,17 +70,6 @@ def _decode(segments, launches=_SPLIT_K, **overrides) -> dict: return _spec(**{**base, **overrides}) -def _mixed(requests, context_lens, launches=_PREFILL, **overrides) -> dict: - base = dict(mixed=(requests, context_lens), slots=32, prefill_pair=True, launches=launches) - return _spec(**{**base, **overrides}) - - -def _short(segments, slots, pool_pages) -> dict: - """Decode rows at low positions over a small pool whose page count and table height - pick the classes of the plain-int `num_pages` and `num_requests` Triton specializes on.""" - return _decode(segments, slots=slots, pool_pages=pool_pages) - - @triton.jit def _e2m1_to_f32(code): """e2m1 code (0..15, bit 3 the sign) -> its value, exactly.""" @@ -866,62 +854,15 @@ class QsaPagedSparseGqaNvfp4(OpWrapper): CELLS: tuple[Cell, ...] = ( Cell( - why="decode: 1 row past the budget, 64-way split-K + merge, PDL", - spec=_decode(((4000, 1),)), - ), - Cell( - why="decode MTP draft 7: 16 requests x 8 rows (128 rows, 8-way)", - spec=_decode(tuple((100 + (499 * g) % 8000, 8) for g in range(16))), - ), - Cell( - why="300 rows, no prefill pair: split-K, one split, no merge", - spec=_decode(((600, 1),) * 300, launches=("splitk",), slots=384), - ), - Cell( - why="context 128 rows from 2032: pair passed, still split-K", - spec=_decode(((2032, 128),), prefill_pair=True), - ), - Cell( - why="context 4096 rows from 0: rows see 1..4096 tokens", - spec=_decode(((0, 4096),), launches=_PREFILL, prefill_pair=True), - ), - Cell( - why="second-level scales g_k 0.37, g_v 2.5: the dequantize rounds to bf16", - spec=_decode(tuple((4000 + 37 * i, 2) for i in range(4)), global_k=0.37, global_v=2.5), - ), - Cell( - why="holes: all-padding row, foreign request, unmapped page", - spec=_decode(((900, 1),) * 4, holes=True), - ), - Cell( - why="the engine's trunk layout: layer 11's planes among the 24 role pages of a slot", + why="decode: 4 requests x 4 rows, 8-way split-K + merge, in the engine's trunk " + "layout: layer 11's planes among the 24 role pages of a slot", spec=_decode( tuple((4000 + 37 * i, 4) for i in range(4)), slot_pages=24, layer=11, pool_pages=600 ), ), Cell( - why="the trunk layout, context 272 rows from 2032: the prefill kernel", - spec=_decode( - ((2032, 272),), - launches=_PREFILL, - prefill_pair=True, - slot_pages=24, - layer=11, - pool_pages=600, - ), - ), - # Triton compiles once per class of each plain-int argument -- equal to 1, a multiple - # of 16, or anything else: the pool's page count, the table's height (32 and 384 - # above) and the split-K kernel's row count (1, 16, 128 and 300 above). The two below - # add a pool of a multiple of 16 pages and table heights 1 and 8, so each class of - # each argument is reached, not every combination of them. - Cell( - why="specialization: pool of 100 pages (other), table height 1", - spec=_short(((600, 1),), slots=1, pool_pages=100), - ), - Cell( - why="specialization: pool of 96 pages (multiple of 16), table height 8", - spec=_short(((600, 1),) * 4, slots=8, pool_pages=96), + why="context 4096 rows from 0: rows see 1..4096 tokens", + spec=_decode(((0, 4096),), launches=_PREFILL, prefill_pair=True), ), Cell( why="captured decode: G 1, 8 and 32 requests x 1 and 8 rows, replay == eager", @@ -949,9 +890,9 @@ class QsaPagedSparseGqaNvfp4(OpWrapper): The decode site reaches 256 rows at max_draft_len 7 (32 requests x 8). Equivalence: the test also runs QSA's bf16 sparse GQA over a bf16 pool - holding the dequantized values; on every cell the two outputs agree within - the test's band. They are not bitwise equal: the dequantizing load changes - the compiled kernel around the same arithmetic. + holding the dequantized values; on every cell but the captured one the two + outputs agree within the test's band. They are not bitwise equal: the + dequantizing load changes the compiled kernel around the same arithmetic. CUDA graphs: the decode site captures at G requests x 1..8 rows over the [32, 260] table; a replay after new operands are written in place is @@ -960,8 +901,8 @@ class QsaPagedSparseGqaNvfp4(OpWrapper): Compiled variants: the kernels take the pool's page count (`num_pages`), the table's row count (`num_requests`) and the split-K kernel the step's row count as plain ints, and Triton compiles each apart at 1, at multiples of 16 - and otherwise; the cells reach each class of each of the three, not every - combination. Every other plain int -- the strides, the table width 260 and + and otherwise. Every cell runs a 32-row table (a multiple of 16); the other + table-height classes are not certified. Every other plain int -- the strides, the table width 260 and the selection's row stride 2051 -- is fixed at the served value the cells use. Preconditions: @@ -973,14 +914,18 @@ class QsaPagedSparseGqaNvfp4(OpWrapper): - `D` a power of two >= 32; `Hq % Hkv == 0`; `tokens_per_block % 4 == 0`. The callable rejects loudly (ValueError) every operand whose dtype, shape or - layout it cannot read; `rows == 0` returns an empty `[0, Hq, D]` tensor. - Calls are bitwise deterministic. + layout it cannot read (the test pins the half-passed prefill pair); `rows == + 0` returns an empty `[0, Hq, D]` tensor (not certified). Calls are bitwise + deterministic. Every row, request, page and position index is widened to 64 bits before any product. Not certified: head geometries other than 24/2 x 256; `tokens_per_block` other than 32; more than 8192 rows; capture of the - context/mixed site; a pool of one slot; an fp16 or unsliced q; int64 request - indices. + context/mixed site; a pool of one slot; second-level scales other than 1.0; + the prefill pair on the split-K kernel (256 context rows or fewer); one split + without a merge; the trunk layout on the prefill kernel; rows whose selection + resolves no token (all padding, a foreign request, unmapped pages); an fp16 + or unsliced q; int64 request indices. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_prefill_compress.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_prefill_compress.py index d70f62d41d8a..415956a05579 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_prefill_compress.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_prefill_compress.py @@ -15,7 +15,8 @@ from tensorrt_llm._torch.attention.backends.sparse.qsa import kernels as _qsa_kernels #: The mRoPE sections the checkpoint declares; the target passes `None` today -#: (`_QSA_MROPE_SECTION`) and these are certified for when it flips. +#: (`_QSA_MROPE_SECTION`) and these are certified, on the NVFP4 trunk-plane cell, +#: for when it flips. MROPE_SECTION = (11, 11, 10) #: This checkpoint's index geometry and the engine's QSA metadata: index head @@ -31,13 +32,9 @@ mrope_section=None, logical_dtype=torch.int64, request_dtype=torch.int32, - distinct_axes=True, layout="plain", - duplicate=False, ) -_DECODE = ((4095, 1), (4096, 1), (4097, 1), (4098, 1), (4099, 1), (8191, 1), (3, 1), (2047, 1)) - #: 32 requests of one MTP7 verify block each (8 rows, two completed groups), at starts that put #: some blocks across a page edge and some inside one. _MTP7_32 = tuple((29 + 61 * i, 8) for i in range(32)) @@ -48,8 +45,8 @@ def _spec(segments, pages, **overrides) -> dict: return {**_CKPT, "segments": segments, "pages": pages, **overrides} -def _old(segments, pages, **overrides) -> dict: - """The old receipt's scenes: prefill from 0 under mRoPE (11, 11, 10).""" +def _mrope(segments, pages, **overrides) -> dict: + """Prefill from 0 under mRoPE (11, 11, 10).""" return _spec(segments, pages, mrope_section=MROPE_SECTION, **overrides) @@ -87,34 +84,20 @@ class QsaPrefillCompress(OpWrapper): ARCHS = frozenset({Arch.SM_103}) CELLS: tuple[Cell, ...] = ( - # First on purpose: the process's first, JIT-compiling call. - Cell( - why="prefill 2 x 128 from 0, mRoPE, distinct t/h/w: cold call", - spec=_old(((0, 128),) * 2, pages=8), - ), Cell( why="prefill 2 x 128 from 0, mrope_section None: the target's call", spec=_spec(((0, 128),) * 2, pages=8), ), - Cell( - why="decode: 8 rows to 8191, anchors and non-anchors, None", - spec=_spec(_DECODE, pages=1024), - ), Cell( why="decode MTP7: 32 requests x 8 rows, two anchors each, None", spec=_spec(_MTP7_32, pages=1100, slots=40, rounding_flips=1), ), # Under an NVFP4 KV pool the position page shares each physical page with the index # pages, so both caches are strided views of one allocation: index page stride - # 12 x 4096 + 192 = 49344 on the trunk (4288 on the draft), position strides - # (24672, 3, 1) (2144, 3, 1). + # 12 x 4096 + 192 = 49344 on the trunk, position strides (24672, 3, 1). Cell( why="NVFP4 trunk planes: 12 index roles and the position page share each page", - spec=_old(((0, 64),) * 2, pages=4, layout="shared_page", roles=12), - ), - Cell( - why="NVFP4 draft planes: one index role and the position page share each page", - spec=_old(((0, 64),) * 2, pages=4, layout="shared_page", roles=1), + spec=_mrope(((0, 64),) * 2, pages=4, layout="shared_page", roles=12), ), ) @@ -138,7 +121,7 @@ class QsaPrefillCompress(OpWrapper): dim of `cos_sin` are read with an assumed unit stride while the callable checks only shape: each was measured accepted and wrong (max abs 4.72, 6.20, 0.80, 7.16). `is_valid` refuses them. - The callable rejects loudly: metadata shape/dtype mismatches, a non-int32 + The callable rejects loudly (not pinned by a test): metadata shape/dtype mismatches, a non-int32 or non-2-D table, more than one index head, a position plane with fewer than 3 axes or not int32, a non-bf16 cache, page sizes that disagree with `tokens_per_block`, a malformed cos|sin table, `4 R != D`, a non-2^k `D` @@ -148,14 +131,15 @@ class QsaPrefillCompress(OpWrapper): `cos_sin`'s dtype is the rotation's precision: bf16 is accepted and was measured 1.4% of elements off by up to 1.56e-2 -- not certified. Text - rows reduce mRoPE to plain RoPE bit for bit. Plain `@triton.jit`, no + rows reduce mRoPE to plain RoPE bit for bit (measured, not pinned by a + test). Plain `@triton.jit`, no autotune; CUDA-graph capturable at a fixed row count (captured at 256 rows, 32 MTP7 verify blocks). Under an NVFP4 KV pool both caches are strided views of one per-slot allocation (each physical page holds the index roles, then the position - page); the cells check the other index roles and the position page - untouched. + page); the trunk-plane cell checks the other index roles and the position + page untouched. The kernel's `tl.rsqrt` is the hardware approximation, an fp32 ulp from the reference's: every cell is bitwise the reference but the 32-block MTP7 one, where @@ -163,11 +147,12 @@ class QsaPrefillCompress(OpWrapper): Hard limits (read from the kernel, not probed): positions, requests and page ids are widened to int64 before any address arithmetic, so no int32 - boundary inside the call. Certified extent: positions to 8191 through a - [32, 260] table over 1024 pages. + boundary inside the call. Certified extent: positions to 1928 through a + [40, 260] table over 1100 pages. Not certified: `compress_ratio` other than 4, `D` other than 128, a bf16 - or fp16 `cos_sin`, positions past the table's rows, multiple devices, + or fp16 `cos_sin`, positions past 1928 or the table's rows, mRoPE outside + the NVFP4 trunk-plane cell, the NVFP4 draft planes (one index role a page), multiple devices, `tokens_per_block` other than 32 (no group straddles a page), eps other than 1e-6, int32 positions or int64 request indices, a strided view of the index cache, position plane or table other than the NVFP4 pool's shared diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py index 1136b8dcee13..86a4f4a92b8f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py @@ -683,14 +683,6 @@ class QsaSparseGqaRowtile(OpWrapper): why="con 128 mixed: 3 chunks (16240 rows) + 52 decodes on the lone pass", spec=_spec([(0, 8192), (2192, 6000), (6144, 2048)] + _decodes(52), slots=128), ), - Cell( - why="int64 request indices and positions", - spec=_spec([(1000, 300)], request_dtype=torch.int64), - ), - Cell( - why="holes: all-padding row, foreign request, unmapped page", - spec=_spec([(900, 40), (950, 40)], holes=True), - ), ) note = """ @@ -736,6 +728,10 @@ class QsaSparseGqaRowtile(OpWrapper): prefill path, which runs eager. Results are deterministic: the bitmap OR, the re-indexed bitmap OR and the lone list order do not change any row's arithmetic. + + Not certified: int64 request indices or positions; rows whose selection + resolves no token (all padding, a foreign request, unmapped pages); + `rows == 0`. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_expand_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_expand_decode.py index 470b7d0dab18..542d4d9a4e70 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_expand_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_expand_decode.py @@ -497,17 +497,17 @@ class QsaTopkExpandDecode(OpWrapper): ARCHS = frozenset({Arch.SM_103}) CELLS: tuple[Cell, ...] = ( - Cell( - why="decode, 1 row at position 8700: 2175 visible blocks of 2560", - spec=dict(rows=1, pos=(8700,)), - ), Cell( why="decode, 4 rows across the budget, one below it", spec=dict(rows=4, pos=(8700, 4095, 1500, 9900)), ), + # The target writes a context chunk's rows through `out` into its rows of the + # step's selection; relu-sparse scores put the zero class (negative zeros among it) + # across the budget, the ties the canonical rule decides. Cell( - why="decode, 1 row with ties at the threshold", - spec=dict(rows=1, pos=(8700,), ties=True), + why="context chunk: 2048 rows of one request (positions 6144..8191), relu-sparse " + "scores, written through `out` into rows of a step-sized selection", + spec=dict(rows=2048, chunk=(6144, 2048), relu_sparse=True, out_margin=64), ), Cell( why="decode, 1 row of scores within 5% of each other", @@ -525,7 +525,12 @@ class QsaTopkExpandDecode(OpWrapper): Compiled once per score-table width and row stride (the block table's, fixed for an engine), k, compress ratio, token budget, index dtypes and PDL: the row count is a - runtime argument, so an eager step's context rows, any count, reuse the kernel. + runtime argument, so an eager step's context rows, any count, reuse the kernel (read + from the launch, not pinned by a test). + + Certified: decode batches of 1 and 4 rows over a 2560-wide score table, a 2048-row + context chunk written through `out`, and a replay under a CUDA graph that reads the + scores after its wait. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_select.py index 295bfa6d32a2..fa7e00d6e456 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_select.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_select.py @@ -74,17 +74,6 @@ class QsaTopkSelect(OpWrapper): # 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 MTP7: 32 requests x 8 rows = 256 rows, some below the budget, NaN tail", spec=dict( @@ -132,15 +121,17 @@ class QsaTopkSelect(OpWrapper): reads out of bounds (not probed). - Logits inside the window are finite. `-inf` is never selected, but a NaN ranks above every finite value and IS selected: NaN in columns 0-3 - of a 32-wide window with `k=8` returned all four. Outside the window - anything goes -- the cells poison it with NaN. + of a 32-wide window with `k=8` returned all four (measured, not pinned + by a test). Outside the window anything goes -- the cells poison it + with NaN. - `index_topk` is the BLOCK budget (512 here). The default, 2048, is the token budget, and with a 512-wide buffer it is the layout fault `is_valid` refuses. - Loud rejections (measured): bf16/fp16 `logits` ("expected scalar type - Float"), int64 `indices` ("expected scalar type Int"), `row_starts` / - `row_ends` shorter than the row count ("one entry per row"). + Loud rejections (measured, not pinned by a test): bf16/fp16 `logits` + ("expected scalar type Float"), int64 `indices` ("expected scalar type + Int"), `row_starts` / `row_ends` shorter than the row count ("one entry per + row"). The CUDA op dispatches by row index, not column count: rows below 12288 run an insertion-sort variant, rows at or above a radix variant, in one launch. @@ -153,9 +144,9 @@ class QsaTopkSelect(OpWrapper): 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]` and `[256, 2048]` (32 requests at - max_draft_len 7), `k=512`, and replayed against new logits and a new - `row_ends` through the captured addresses. Bitwise comparison of two + CUDA graphs: captured at `[256, 2048]` (32 requests at max_draft_len 7), + `k=512`, and replayed against new logits and a new `row_ends` through the + captured addresses. Bitwise comparison of two replays is not a valid check (order is unspecified). Hard limits: not probed. `indices` is addressed as `row * index_topk + diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_qk_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_qk_norm_rope.py index 8613a41e3c88..27a17ed46ace 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_qk_norm_rope.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_qk_norm_rope.py @@ -975,6 +975,8 @@ class Qwen4ExpQkNormRope(OpWrapper): ARCHS = frozenset({Arch.SM_103}) + # Four cells, one per compiled launch the forward reaches: the Triton kernel's 1-, 2- + # and 8-head launches (by decode batch size) and the CuTe kernel (prefill). CELLS: tuple[Cell, ...] = ( Cell( why="decode M=1: one token on the 1-head/2-warp launch, three distinct mRoPE rows", @@ -985,14 +987,10 @@ class Qwen4ExpQkNormRope(OpWrapper): spec=dict(tokens=17, **CKPT), ), Cell( - why="decode M=65: first on the 8-head launch (last main program half masked)", - spec=dict(tokens=65, **CKPT), - ), # fmt: skip - Cell( - why="decode M=128: max concurrency, positions up to the table end (262144 rows)", + why="decode M=128: max concurrency, 8-head launch (last main program half masked), " + "positions up to the table end (262144 rows)", spec=dict(tokens=128, max_position=262144, **CKPT), ), - Cell(why="prefill 1024 tokens: first on the CuTe kernel", spec=dict(tokens=1024, **CKPT)), Cell(why="prefill 16384 tokens: max_num_tokens", spec=dict(tokens=16384, **CKPT)), ) @@ -1015,12 +1013,12 @@ class Qwen4ExpQkNormRope(OpWrapper): bounded by load latency at three CTAs an SM). Numerics: fp32 throughout, one rounding at the store; the CuTe kernel is bitwise - the Triton kernel's eight-head launch (the test checks it at 1024-16384 tokens, - signed zeros included), so the cells' bands cover both. Against the fp64 + the Triton kernel's eight-head launch (measured at 1024-16384 tokens, signed zeros + included; not pinned by a test). Against the fp64 reference over the same table every cell is within one output rounding step (measured max 1.48 ulp of the row scale, relative RMS <= 0.011 ulp). - Against the CUDA pair it replaces, <= 1.41 ulp of the row scale and 0.09 - ulp RMS (that op derives cos/sin with `__sincosf`, whose phase error grows + Against the CUDA pair it replaces (measured, not pinned by a test), <= 1.41 + ulp of the row scale and 0.09 ulp RMS (that op derives cos/sin with `__sincosf`, whose phase error grows with position). Three agreeing position rows select the same table row for every pair, so a text batch is exactly the plain rope on row 0. @@ -1032,7 +1030,11 @@ class Qwen4ExpQkNormRope(OpWrapper): Hard limits: token offsets in int64; positions index the table in int64. - Not certified: fp16/fp32 activations, `rotary_dim` > either head width, + CUDA graphs: the Triton kernel captured at M=4 replays bitwise an eager call on + inputs rewritten in place. + + Not certified: non-zero index-plane mRoPE sections (the target passes (0, 0)), + capture of the CuTe kernel, `T == 0`, fp16/fp32 activations, `rotary_dim` > either head width, non-power-of-two head widths or head counts (tl.arange), YaRN, non-Gemma norm, a k plane on the index side (its scratch slot is not written). """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/vision_rope_2d.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/vision_rope_2d.py index 90917dac79e2..6e5d284f768b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/vision_rope_2d.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/vision_rope_2d.py @@ -108,20 +108,6 @@ class VisionRope2d(OpWrapper): grid_dtype=torch.int64, ), ), - Cell( - why="grid side 6272 (2x6272) float32: largest width angle, fp32 exposes the table", - spec=dict(grid=((1, 2, 6272),), dtype=torch.float32), - ), - Cell( - why="grid side 6272 (6272x2) bfloat16: largest height angle", - spec=dict(grid=((1, 6272, 2),), dtype=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)) ) note = """ @@ -130,23 +116,25 @@ class VisionRope2d(OpWrapper): `k` need last-dim stride 1 only; they are not mutated. `head_dim` must be divisible by 4; otherwise the angle table's width - mismatches the head and the broadcast raises (loud, so not guarded). + mismatches the head and the broadcast raises (loud, so not guarded; not + pinned by a test). - Accuracy vs fp64 angles over the cells: fp32 <= 4.77e-7, bf16 <= 1.56e-2 - (one ulp at unit-normal operands), flat in the grid side up to 6272 -- the - fp64 table is why (an fp32 `inv_freq` times fp32 position - moved an fp32 output 1.475e-5 at a 64 side, outside the default gate). - Wrong operators at 1024 bf16 tokens land 8.5-9.5 away: GPT-J pairing, axes - swapped, raster order, `inv_freq` over `head_dim`. + Accuracy vs fp64 angles over the cells: bf16 <= 1.56e-2 (one ulp at + unit-normal operands). Measured, not certified: fp32 <= 4.77e-7, flat in the + grid side up to 6272 -- the fp64 table is why (an fp32 `inv_freq` times fp32 + position moved an fp32 output 1.475e-5 at a 64 side, outside the default + gate); wrong operators at 1024 bf16 tokens land 8.5-9.5 away: GPT-J pairing, + axes swapped, raster order, `inv_freq` over `head_dim`. Host-bound: ~353 us/call at 1024 and 4096 tokens (`.tolist()` + small - launches), x27 blocks per image. The table depends only on the grid. - CUDA graphs: capture bakes in the grid; replay under another grid was not - measured. - - Not certified: float16; mixed q/k dtypes or head counts; float8; - `spatial_merge_size` != 2; `theta` other than 1e4; grid sides past 6272 or - calls past 32768 rows. + launches), x27 blocks per image. The table depends only on the grid, and the + result does not depend on the grid's placement (host or device) or integer + dtype (int32 or int64). + + Not certified: float16 or float32; mixed q/k dtypes or head counts; head + geometries other than 16 x 72; float8; `spatial_merge_size` != 2; `theta` + other than 1e4; grid sides past 90 or calls past 32768 rows; CUDA-graph + capture (it would bake in the grid). """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/embedding/ngram_hash.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/embedding/ngram_hash.py index 1892ee2c7cac..e3841c25f537 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/embedding/ngram_hash.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/embedding/ngram_hash.py @@ -76,7 +76,6 @@ class NgramHash(OpWrapper): # token from `ngram_token_history`, num_query_tokens 1. CELLS: tuple[Cell, ...] = ( _rows(4, 3, 1, "decode row whose cached context is EOS", eos_positions=(0, 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 ), @@ -97,7 +96,7 @@ class NgramHash(OpWrapper): gives the identical result). Certified on an int64 `token_history`, the window `ngram_token_history` returns. Negative ids are accepted and hashed as ordinary ids, not padding. Longer histories and `num_query_tokens > 1` - (a chunk's tail) are exercised by the test beyond the cells, not certified. + (a chunk's tail) are not certified. A single multiplier raises `ZeroDivisionError`. The four silent mistakes are `is_valid`'s. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/bf16_gemv_pdl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/bf16_gemv_pdl.py index a70d11d4f21e..c4b93e549701 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/bf16_gemv_pdl.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/bf16_gemv_pdl.py @@ -323,14 +323,13 @@ class Bf16GemvPdl(OpWrapper): ARCHS = frozenset({Arch.SM_103}) # Every shape is a qwen3.8-flash-next-nvfp4 tp1 decode projection; the target calls - # it at one row (`_decode_projection`). + # it at one row (`_decode_projection`). Four cells: each `TILING` key is its own + # compiled kernel the decode forward reaches, plus the N=1 edge. CELLS: tuple[Cell, ...] = ( - _cell("qwen3.8 decode M=1 router N=512", 1, 512, 2560), - _cell("qwen3.8 decode M=1 shared gate N=1, outputs per CTA halved to 1", 1, 1, 2560), - _cell("qwen3.8 decode M=1 shared down K=640", 1, 2560, 640), _cell("qwen3.8 decode M=1 GDN out / attention o, K=6144 shared-memory tier", 1, 2560, 6144), _cell("qwen3.8 decode M=1 GDN in-proj concat N=16480, L2 prefetch grid", 1, 16480, 2560), - _cell("qwen3.8 decode M=4 (MAX_ROWS) GDN out / attention o", 4, 2560, 6144), + _cell("qwen3.8 decode M=1 shared down K=640", 1, 2560, 640), + _cell("qwen3.8 decode M=1 shared gate N=1, outputs per CTA halved to 1", 1, 1, 2560), ) note = """ @@ -352,12 +351,15 @@ class Bf16GemvPdl(OpWrapper): order differs from the fp32 reference. Hard limits - M rows per launch boundary MAX_ROWS=4 driven at 1 and 4 (cells) + M rows per launch boundary MAX_ROWS=4 driven at 1 (cells) K a TILING key 2560, 6144, 640 all driven N any; the outputs per CTA halve until they divide N N=1 driven - PDL off (`pdl=False`, or PDL disabled in the environment) is the same kernel - without the attribute; the test holds it bitwise to PDL on. + The test also replays a PDL chain of these GEMVs under a CUDA graph, each + reading the one before it. + + Not certified: M = 2..4; PDL off (`pdl=False`, or PDL disabled in the + environment), the same kernel without the attribute. """ def __call__(self, mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool = True) -> torch.Tensor: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/low_m_bf16_mm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/low_m_bf16_mm.py index 8405f6fb8ca4..967a23116a87 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/low_m_bf16_mm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/low_m_bf16_mm.py @@ -145,7 +145,8 @@ class LowMBf16Mm(OpWrapper): ARCHS = frozenset({Arch.SM_103}) # Every cell is a qwen3.8-flash-next-nvfp4 tp1 decode projection at M >= 2 (one row - # takes `bf16_gemv_pdl` in the target); `kind` pins the branch `plan` takes for it. + # takes `bf16_gemv_pdl` in the target), one per branch; `kind` pins the branch `plan` + # takes for it. The test's host-only plan pins hold the table's other shapes. CELLS: tuple[Cell, ...] = ( _cell( "qwen3.8 decode M=2 GDN in_z: direct, 2-row tile, pdl on", @@ -156,16 +157,13 @@ class LowMBf16Mm(OpWrapper): pdl=True, ), _cell( - "qwen3.8 decode M=8 o_proj: split-K 3 over the cluster, pdl on", - 8, + "qwen3.8 decode M=32 o_proj: split-K, mma_n 32, pdl on (M boundary)", + 32, 2560, 6144, "splitk", pdl=True, ), - _cell("qwen3.8 decode M=32 GDN in_a/in_b: direct, 8-row tiles x4", 32, 48, 2560, "direct"), - _cell("qwen3.8 decode M=32 o_proj: split-K, mma_n 32", 32, 2560, 6144, "splitk"), - _cell("qwen3.8 decode M=2 GDN in-proj concat: cuBLAS branch", 2, 16480, 2560, "cublas"), _cell("qwen3.8 decode M=128: cuBLAS branch above M=32", 128, 10240, 2560, "cublas"), ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_gated_residual.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_gated_residual.py index 38b3a8de5452..9983a66b1c55 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_gated_residual.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_gated_residual.py @@ -71,16 +71,15 @@ class Qwen4ExpGatedResidual(OpWrapper): # spec: leading (token dims), inject (True). Every cell is this # checkpoint's 4x2560 bf16 geometry with synthetic weights of its measured # element statistics, the normed stream reaching ~+-30 as the real one does. - # Leading dims beyond one are `test_flattened_call_matches_batched`'s. - CELLS: tuple[Cell, ...] = tuple( - Cell(why=f"checkpoint 4x2560 bf16, {why}", spec=dict(leading=(t,))) - for t, why in ( - (1, "decode, 1 token"), - (8, "decode batch 8 / MTP draft rows"), - (512, "prefill 512 tokens"), - (8192, "prefill 8192 tokens, largest measured"), - ) - ) + ( + CELLS: tuple[Cell, ...] = ( + Cell( + why="checkpoint 4x2560 bf16, decode batch 8 / MTP draft rows", + spec=dict(leading=(8,)), + ), + Cell( + why="checkpoint 4x2560 bf16, prefill 8192 tokens, largest measured", + spec=dict(leading=(8192,)), + ), Cell( why="stack-final mixer: block_inject_weight None", spec=dict(leading=(256,), inject=False), @@ -95,7 +94,7 @@ class Qwen4ExpGatedResidual(OpWrapper): blocks and `shared_head`'s own mixer). Preconditions, all loud (each measured to raise `RuntimeError` before - producing anything): the width divisible by `hc_count` (unflatten), the + producing anything; the width one is tested): the width divisible by `hc_count` (unflatten), the weights' input width equal to it and their dtype equal to the activation's (`F.linear`). @@ -111,13 +110,16 @@ class Qwen4ExpGatedResidual(OpWrapper): sigmoid saturates (a property of the function, outside the checkpoint's range). - CUDA-graph capture certified, bitwise equal to eager. Cost, not gated, - shared device: ~235 us at 1..512 tokens (overhead-bound), 475 us at - 4096, 902 us at 8192. + Certified: bf16 2-D `[rows, 10240]` at 8 and 8192 rows with the + injection, 256 rows without it (the stack-final mixer); CUDA-graph replay + at 8 rows, bitwise equal to eager. Cost, not gated, shared device: ~235 us + at 1..512 tokens (overhead-bound), 475 us at 4096, 902 us at 8192. - Not certified: hc_count other than 4, hidden other than 2560, lowrank - other than 320, float16 and float32, weight/activation dtype mixing. Hard - limits: not probed (largest operand 8192 x 10240 = 2^26.3 elements). + Not certified: other row counts, more than one leading dim, hc_count other + than 4, hidden other than 2560, lowrank other than 320, float16 and + float32, weight/activation dtype mixing; the operand-range bounds above + are measured. Hard limits: not probed (largest operand 8192 x 10240 = + 2^26.3 elements). """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first_cute.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first_cute.py index 656d25581046..8e58e9df35cc 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first_cute.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first_cute.py @@ -428,17 +428,18 @@ class Qwen4ExpHcDecodeFirstCute(OpWrapper): # (0: top-k of the finalize folded in, block_out then the permuted rows), drop # (False: some slots -1), n_cols (324). The weight is the target's padded # [width, 384] (columns past n_cols never read). + # Four cells: the forward reaches four compiled variants (inject at 324 cols, the + # fold at 324, the fold at 320 into the stack-final mixer, norm-only), none of which + # another covers; the -1 slots ride the 324-col fold cell. CELLS: tuple[Cell, ...] = ( Cell(why="checkpoint 4x2560, 324 cols of 384, bf16, inject, 1 row", spec=dict(rows=1)), - Cell(why="MoE combine, [T, 1] gate, 1 row", spec=dict(rows=1, moe=True)), - Cell(why="finalize fold, top-10, 1 row (con 1)", spec=dict(rows=1, moe=True, fold=10)), Cell( - why="finalize fold into the stack-final mixer: 320 cols", - spec=dict(rows=1, moe=True, fold=10, n_cols=320), + why="finalize fold, top-10, last 3 slots mapped to -1, 1 row", + spec=dict(rows=1, moe=True, fold=10, drop=True), ), Cell( - why="finalize fold, last 3 slots mapped to -1", - spec=dict(rows=1, moe=True, fold=10, drop=True), + why="finalize fold into the stack-final mixer: 320 cols", + spec=dict(rows=1, moe=True, fold=10, n_cols=320), ), Cell(why="norm-only site (stack start / after PLE)", spec=dict(rows=1, inject=False)), ) @@ -462,15 +463,20 @@ class Qwen4ExpHcDecodeFirstCute(OpWrapper): finalize and the combine likewise); `part` and `ss` are fp32 reductions in a fixed order, inside the float64 reference's band (the combine and the finalize can each land one ulp of the block output from it, scaled into `hc_out` by `inject`). - Deterministic. Graph-capturable: compiled on the first call per variant, which - must precede capture (the engine's warmup forward does it). + Deterministic (measured). Graph-capturable: compiled on the first call per + variant, which must precede capture (the engine's warmup forward does it). PDL: with `TRTLLM_ENABLE_PDL` (default on, sm_90+) the weight rows and the norm slice stream before `griddepcontrol.wait`, every activation read and every write after it, and dependents are released before the wait. - Not certified: a flat `[T]` gate, a row-slice `hyper_input`, a weight - narrower than the padded 384 rows (accepted; no cell drives them). + Certified: one bf16 row, the injection at 324 columns, the fold (top-10, three + slots -1) at 324 columns, the fold into the stack-final mixer at 320 columns, + the norm-only site; graph replay under PDL with the fold. + + Not certified: the combine without the fold, a flat `[T]` gate, a row-slice + `hyper_input`, a weight narrower than the padded 384 rows (accepted; no cell + drives them). """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix_fp4_cute.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix_fp4_cute.py index 30cd06f39156..45a3c08057bd 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix_fp4_cute.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix_fp4_cute.py @@ -368,11 +368,8 @@ def _reference(part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject): return mixed, inject_out -_CELLS = ( - Cell(why="checkpoint 4x2560 r320 bf16, inject, 1 row (con 1)", spec=dict(rows=1)), - Cell(why="stack-final mixer (no inject), 1 row", spec=dict(rows=1, with_inject=False)), - Cell(why="hc_out a row slice buf[::2], 1 row", spec=dict(rows=1, layout="row_slice")), -) +_INJECT = Cell(why="checkpoint 4x2560 r320 bf16, inject, 1 row (con 1)", spec=dict(rows=1)) +_ROW_SLICE = Cell(why="hc_out a row slice buf[::2], 1 row", spec=dict(rows=1, layout="row_slice")) class Qwen4ExpHcDecodeMixFp4Cute(OpWrapper): @@ -400,7 +397,8 @@ class Qwen4ExpHcDecodeMixFp4Cute(OpWrapper): """ ARCHS = frozenset({Arch.SM_103}) - CELLS: tuple[Cell, ...] = _CELLS + # Its call sites always inject, so no stack-final (no-inject) cell. + CELLS: tuple[Cell, ...] = (_INJECT, _ROW_SLICE) note = """ Metadata consumed: none. Call sites: `DecodeTarget`'s MoE sites at a one-row @@ -414,8 +412,12 @@ class Qwen4ExpHcDecodeMixFp4Cute(OpWrapper): Numerics: `mixed` and `inject_out` as `qwen4_exp_hc_decode_mix_cute`; the NVFP4 is `fp4_quantize`'s arithmetic on the rounded `mixed` (`rcp.approx.ftz` reciprocals, e4m3 scale rounded to nearest, e2m1 nearest-even with saturation at 6). - Deterministic; graph-capturable once compiled (the first call per variant). - PDL: the weight rows and norm elements stream before `griddepcontrol.wait`. + Deterministic (measured); graph-capturable once compiled (the first call per + variant). PDL: the weight rows and norm elements stream before + `griddepcontrol.wait`. + + Certified: one bf16 row with the injection, `hc_out` compact and a row slice; + graph replay under PDL. Not certified: `with_inject=False`. """ def __call__(self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject, global_scale): @@ -460,7 +462,11 @@ class Qwen4ExpHcDecodeMixCute(OpWrapper): """ ARCHS = frozenset({Arch.SM_103}) - CELLS: tuple[Cell, ...] = _CELLS + CELLS: tuple[Cell, ...] = ( + _INJECT, + Cell(why="stack-final mixer (no inject), 1 row", spec=dict(rows=1, with_inject=False)), + _ROW_SLICE, + ) note = """ Metadata consumed: none. Call sites: `DecodeTarget`'s non-MoE sites at a one-row @@ -475,8 +481,12 @@ class Qwen4ExpHcDecodeMixCute(OpWrapper): `up` the `.t()` view of a contiguous `[width, R]` weight. Numerics: fp32 inside (`low`, the gate logits and `normed` are never rounded), one - rounding per output, inside the float64 reference's band. Deterministic; - graph-capturable once compiled (the first call per variant). + rounding per output, inside the float64 reference's band. Deterministic + (measured); graph-capturable once compiled (the first call per variant). + + Certified: one bf16 row with and without the injection, `hc_out` compact and a + row slice. Not certified: graph replay of this build (the NVFP4 build's + replay test covers the shared kernel, not this compiled variant). """ def __call__(self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject): diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_tc_first.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_tc_first.py index 7de60f75d3e6..df942de81ed0 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_tc_first.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_tc_first.py @@ -564,7 +564,6 @@ class Qwen4ExpHcDecodeTcFirst(OpWrapper): # spec: rows, inject (True). The checkpoint's geometry throughout: hc_count 4, # rank 320, hidden 2560, the padded [384, K] weight, bf16. CELLS: tuple[Cell, ...] = ( - Cell(why="checkpoint 4x2560, 324 cols of 384, inject, decode 2 rows", spec=dict(rows=2)), Cell( why="checkpoint 4x2560, 324 cols of 384, inject, decode 128 rows (MAX_ROWS)", spec=dict(rows=128), @@ -573,7 +572,7 @@ class Qwen4ExpHcDecodeTcFirst(OpWrapper): why="stack-final mixer: 320 cols, no inject, 128 rows", spec=dict(rows=128, inject=False), ), - Cell(why="zero rows", spec=dict(rows=0)), + Cell(why="checkpoint 4x2560, 324 cols of 384, inject, decode 2 rows", spec=dict(rows=2)), ) note = """ @@ -604,8 +603,11 @@ class Qwen4ExpHcDecodeTcFirst(OpWrapper): against 8.6 us for cuBLAS split-K + reduce + lowrank_act; tile sweep BN 32 / 48 / 64 / 112 x KS 8 / 10 / 16, 48 x 16 fastest at every row count. - Not certified: hc_count other than 4, an unpadded weight (N = R + C), a - single row (below the dispatch range; accepted, no cell drives them). + Certified: bf16, 324 columns with the injection at 2 and 128 rows, 320 + columns without it at 128 rows; graph replay. + Not certified: hc_count other than 4, an unpadded weight (N = R + C), zero + rows, a single row (below the dispatch range; accepted, no cell drives + them), other row counts without the injection. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_inject_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_inject_norm.py index bfb9b7562991..9ddaa39c7e7f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_inject_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_inject_norm.py @@ -783,30 +783,26 @@ class Qwen4ExpHcInjectNorm(OpWrapper): # spec: rows, inject (True), moe (False: shared + gate_logits [T, 1] too), fold # (0: top-k of the finalize folded in), drop (False: some slots -1). The # checkpoint's geometry throughout: hc_count 4, hidden 2560, bf16, contiguous rows. + # Four cells: the forward reaches four distinct launches -- the Triton kernel's + # eight-warp form (decode), its two-warp form (prefill, here with the combine), + # the CuTe fold kernel (decode MoE sites; the -1 slots ride it) and the norm-only + # variant -- none of which another covers. CELLS: tuple[Cell, ...] = ( Cell( why="checkpoint 4x2560 bf16, inject, decode batch 128 (eight-warp form)", spec=dict(rows=128), ), - Cell( - why="checkpoint 4x2560 bf16, inject, prefill chunk 8192 (two-warp form)", - spec=dict(rows=8192), - ), - Cell( - why="norm only (no inject: stack start / after PLE), decode 128", - spec=dict(rows=128, inject=False), - ), Cell( why="MoE combine, [T, 1] gate, prefill chunk 16384 (the ladder's max_num_tokens)", spec=dict(rows=16384, moe=True), ), Cell( - why="MoE combine + finalize fold (top-10, CuTe kernel), decode batch 128", - spec=dict(rows=128, moe=True, fold=10), + why="MoE combine + finalize fold (top-10, CuTe kernel), three slots -1, decode 128", + spec=dict(rows=128, moe=True, fold=10, drop=True), ), Cell( - why="MoE combine + finalize fold, three slots mapped to -1, 8 rows", - spec=dict(rows=8, moe=True, fold=10, drop=True), + why="norm only (no inject: stack start / after PLE), decode 128", + spec=dict(rows=128, inject=False), ), ) @@ -837,7 +833,8 @@ class Qwen4ExpHcInjectNorm(OpWrapper): `rsqrt`, one rounding. Against the float64 reference the combine can land one ulp of the block output away (its rounding midpoint), which the accumulate carries scaled by `inject`; the test's band states that. - Graph-capturable: no host sync, outputs allocated per call. + Graph-capturable by construction (no host sync, outputs allocated per + call); not certified. Launch: one program per (row, stream), stream-minor in a 1-D grid, so a row's `block_out` / `shared` are read from HBM once and from L2 by the @@ -857,13 +854,21 @@ class Qwen4ExpHcInjectNorm(OpWrapper): a row's two CTAs one cluster, CTA j forming half the block output (the finalize and the combine) into both CTAs' shared memory, then each accumulating and normalizing two streams, in the Triton kernel's layout and - reduction tree -- bitwise its `hc_out` and `normed`. GB300, graph replay, + reduction tree -- bitwise its `hc_out` and `normed` (tested at 511 rows, the + last before the switch). GB300, graph replay, 16 rotating sites, us: 128 rows 7.00 -> 5.97, 64 rows 6.63 -> 5.25, 8 rows 3.59 -> 3.09; it releases dependents at its start as the Triton kernel does. + Certified: bf16, the checkpoint's 4 x 2560: the injection at 128 rows + (eight-warp form), the combine at 16384 rows (two-warp form), the combine + with the fold (three slots -1) at 128 rows (CuTe kernel), norm-only at 128 + rows; the CuTe fold bitwise the Triton kernel at 511 rows. + Not certified: hc_count other than 4, hidden other than 2560, float16 - operands, a flat `[T]` gate, row strides other than the width, zero rows - (the kernels take them; no cell drives them). + operands, a flat `[T]` gate, row strides other than the width, zero rows, + the injection without the combine at prefill sizes, the combine without the + fold at decode sizes, CUDA-graph replay (the kernels take them; no cell + drives them). Hard limits: row offsets are int64 (`program_id(0)` widened); the grid is `rows * hc_count` programs along x (2^31 - 1); the column offset diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_lowrank_act.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_lowrank_act.py index 0dec54e742b6..493ca58301c4 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_lowrank_act.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_lowrank_act.py @@ -68,16 +68,10 @@ class Qwen4ExpHcLowrankAct(OpWrapper): ARCHS = frozenset({Arch.SM_103}) # spec: rows, inject (True), hc_count (4), rank (320), logits_width (384), dtype (bf16). - CELLS: tuple[Cell, ...] = tuple( - Cell(why=f"checkpoint rank 320 + 4, padded to 384, {why}", spec=dict(rows=t)) - for t, why in ( - (1, "1 row"), - (128, "128 rows"), - (8192, "prefill 8192"), - ) - ) + ( + CELLS: tuple[Cell, ...] = ( + Cell(why="checkpoint rank 320 + 4, padded to 384, 128 rows", spec=dict(rows=128)), + Cell(why="checkpoint rank 320 + 4, padded to 384, prefill 8192", spec=dict(rows=8192)), Cell(why="stack-final mixer: no inject, 8192 rows", spec=dict(rows=8192, inject=False)), - Cell(why="zero rows", spec=dict(rows=0)), ) note = """ @@ -86,8 +80,11 @@ class Qwen4ExpHcLowrankAct(OpWrapper): projection. Preconditions: `logits` fp32, unit column stride, at least `rank (+ hc_count)` columns; `hc_count` a power of two. Graph-capturable. - Not certified: hc_count other than 4, rank other than 320, float16 - outputs. + Certified: bf16 outputs from 384-wide fp32 logits, with the injection at 128 + and 8192 rows, without it at 8192 rows; graph replay. + + Not certified: zero rows, hc_count other than 4, rank other than 320, + float16 outputs. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean.py index 60db120b3ac7..755fc8535e08 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean.py @@ -835,14 +835,10 @@ class Qwen4ExpHcUpGateMean(OpWrapper): CELLS: tuple[Cell, ...] = tuple( Cell(why=f"checkpoint 4x2560, rank 320, {why}", spec=dict(rows=t)) for t, why in ( - (1, "1 row, one partial tile"), (128, "decode 128 rows, one full tile"), (16384, "prefill chunk 16384 (max_num_tokens)"), ) - ) + ( - Cell(why="ragged tail, 16320 + 70 rows (mixed step)", spec=dict(rows=16390)), - Cell(why="zero rows", spec=dict(rows=0)), - ) + ) + (Cell(why="ragged tail, 16320 + 70 rows (mixed step)", spec=dict(rows=16390)),) note = """ Tier 6 (CuTe DSL written for this target): a persistent tcgen05 GEMM whose @@ -861,8 +857,12 @@ class Qwen4ExpHcUpGateMean(OpWrapper): that releases dependents early (`qwen4_exp_hc_decode_tc_first`) the weight read overlaps the producer. - Not certified: hc_count other than 4, hidden other than 2560, rank other - than 320 (the kernel compiles for them; no cell drives them). + Certified: bf16 at the checkpoint's 4 x 2560, rank 320: 128 rows (one full + tile), 16384 rows, 16390 rows (a partial last tile); graph replay. + + Not certified: zero rows, fewer rows than a tile, hc_count other than 4, + hidden other than 2560, rank other than 320 (the kernel compiles for them; + no cell drives them). """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean_fp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean_fp4.py index 49f2668d3374..6d4cfb844b64 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean_fp4.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean_fp4.py @@ -52,14 +52,10 @@ class Qwen4ExpHcUpGateMeanFp4(OpWrapper): CELLS: tuple[Cell, ...] = tuple( Cell(why=f"checkpoint 4x2560, rank 320, {why}", spec=dict(rows=t)) for t, why in ( - (2, "decode 2 rows (smallest fused-up batch)"), (128, "decode 128 rows (max_con)"), (16384, "prefill chunk 16384 (max_num_tokens)"), ) - ) + ( - Cell(why="ragged tail, 16320 + 70 rows (mixed step)", spec=dict(rows=16390)), - Cell(why="zero rows", spec=dict(rows=0)), - ) + ) + (Cell(why="ragged tail, 16320 + 70 rows (mixed step)", spec=dict(rows=16390)),) note = """ Tier 6 (CuTe DSL written for this target): `qwen4_exp_hc_up_gate_mean`'s @@ -76,11 +72,18 @@ class Qwen4ExpHcUpGateMeanFp4(OpWrapper): Numerics: the quantization is the trtllm NVFP4 recipe instruction for instruction (`rcp.approx.ftz.f32` reciprocals, `cvt.rn.satfinite` to e4m3 - and to e2m1, every multiply `mul.rn.f32`); the GPU test checks it bitwise - against `torch.ops.trtllm.fp4_quantize` on the returned `mixed`. - - Not certified: hc_count other than 4, hidden other than 2560, rank other - than 320 (the kernel compiles for them; no cell drives them). + and to e2m1, every multiply `mul.rn.f32`); the GPU test checks it against + `fp4_quantize`'s native-torch reference on the returned `mixed` (scales + bitwise, e2m1 codes bitwise outside the reciprocal tie window). Bitwise + `torch.ops.trtllm.fp4_quantize`, ties included, and across tiny and + saturating global scales: measured, not certified. + + Certified: bf16 at the checkpoint's 4 x 2560, rank 320, a calibrated global + scale: 128 rows, 16384 rows, 16390 rows (a partial last tile); graph replay. + + Not certified: zero rows, fewer rows than a tile, hc_count other than 4, + hidden other than 2560, rank other than 320 (the kernel compiles for them; + no cell drives them). """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_post_inject.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_post_inject.py index aec5f3330ad4..d51ddde1e610 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_post_inject.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_post_inject.py @@ -43,23 +43,17 @@ class Qwen4ExpPostInject(OpWrapper): ARCHS = frozenset({Arch.SM_103}) # spec: leading (token dims), hc_count, hidden, dtype; layout - # "contiguous" (default) | "row_slice" (`buf[::2]`) | "col_slice" - # (`buf[:, :width]`), both last-dim stride 1. + # "contiguous" (default) | "col_slice" (`buf[:, :width]`, last-dim stride 1). CELLS: tuple[Cell, ...] = tuple( Cell( why=f"checkpoint 4x2560 bf16, {why}", spec=dict(leading=(t,), hc_count=CKPT_HC_COUNT, hidden=CKPT_HIDDEN), ) for t, why in ( - (1, "decode, 1 token"), (64, "decode batch / MTP rows, 64 tokens"), (8192, "prefill 8192 tokens, largest measured"), ) ) + ( - 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"), @@ -94,8 +88,12 @@ class Qwen4ExpPostInject(OpWrapper): broadcast operands take a non-vectorized path. An in-place form was not measured. - Not certified: hc_count other than 4, hidden other than 2560, float16, - float32 and float8 operands, leading dims beyond one. Hard limits: not + Certified: bf16 2-D operands at the checkpoint's 4 x 2560, 64 and 8192 + rows contiguous, 8 rows with `hyper_input` a column slice; graph replay. + + Not certified: other row counts, a row-slice `hyper_input`, hc_count other + than 4, hidden other than 2560, float16, float32 and float8 operands, + leading dims beyond one. Hard limits: not probed (the largest measured operand is 8192 x 10240 = 2^26.3 elements). """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py index 8f993b017a57..9ba3b8122735 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py @@ -43,16 +43,13 @@ def _ckpt( local (tp1), SwiGLU, no clamp, 16-element e4m3 block scales. `mma_n` is the autotuner tactic's MMA N tile (128 is also the fallback when the - tuner has no entry). `dist` names how the test draws the expert ids, as in the - moe_sort entry's test: `router` = each row's top-10 of uniform scores; `hot` = - every row on the same 10 experts; `tile_edges` = three experts hold exactly - tile-1, tile and tile+1 rows; `round_robin` = row i takes experts (10 i + j) mod - 512. `gsf` sets `global_sf` against the activations' magnitude: `calibrated` = - the activation maximum the test calibrates to maps to the e4m3 maximum, as the - checkpoint's per-layer `1 / down_proj.input_scale` does; `saturating` = 25x - above that, so most block scales clip at 448 and the codes at 6; `subnormal` = - 2^-12 of it, so block scales fall into e4m3 subnormals and zero. `zero_tokens` - makes that many leading tokens all-zero rows (zero block scales). + tuner has no entry). `dist` names how the test draws the expert ids: `router` = + each row's top-10 of uniform scores. `gsf` sets `global_sf` against the + activations' magnitude: `calibrated` = the activation maximum the test + calibrates to maps to the e4m3 maximum, as the checkpoint's per-layer `1 / + down_proj.input_scale` does; `subnormal` = 2^-12 of it, so block scales fall + into e4m3 subnormals and zero. `zero_tokens` makes that many leading tokens + all-zero rows (zero block scales). """ return dict( tokens=tokens, @@ -142,22 +139,10 @@ class CuteDslNvfp4GatherGroupedGemmActFusion(OpWrapper): why="tile 128: decode, 1 request without MTP", spec=_ckpt(1, 128), ), - Cell( - why="tile 128: MMA N 256 tactic, 32 requests at MTP verify width 4; largest captured T", - spec=_ckpt(128, 128, mma_n=256), - ), Cell( why="tile 128: 8192 rows, max_num_tokens prefill", spec=_ckpt(8192, 128), ), - Cell( - why="tile 128: expert rows tile-1, tile, tile+1; the mn_limit edge", - spec=_ckpt(300, 128, dist="tile_edges"), - ), - Cell( - why="tile 128: global_sf 25x calibrated; block scales clip at 448", - spec=_ckpt(2048, 128, gsf="saturating"), - ), Cell( why="tile 128: tiny global_sf and 7 all-zero token rows; subnormal and zero " "block scales", @@ -188,8 +173,8 @@ class CuteDslNvfp4GatherGroupedGemmActFusion(OpWrapper): (tile_size, 256) and with no tuning entry falls back to (tile_size, 128); serving warmup tunes per row bucket and may pick either. Every pair's codes and block scales are bitwise identical across tile 128 / 256 x MMA N 128 / 256 - x the fallback (the test checks it at 128 and 8192 rows), so cells driven at - one tactic certify the others. + x the fallback (the test checks it at 8192 rows), so cells driven at one + tactic certify the others. Determinism. A (token, expert) pair's output row depends on that pair alone: it is bitwise unchanged by reordering an expert's rows (moe_sort orders them by @@ -203,11 +188,12 @@ class CuteDslNvfp4GatherGroupedGemmActFusion(OpWrapper): the same tile arrays, and the combine reads only expanded_idx_to_permuted_idx's rows. - Scale range: with global_sf 25x the calibrated value the block scales clip at - 448 and the codes at 6 (no NaN); with a tiny one the block scales fall to e4m3 - subnormals and to zero. An all-zero token gives scale 0 and codes of magnitude - 0. A block whose scale is 0 dequantizes to zero whatever its codes hold, and - the test's compare does not read them. + Scale range: with a tiny global_sf the block scales fall to e4m3 subnormals + and to zero (a cell). An all-zero token gives scale 0 and codes of magnitude + 0 (the same cell). With global_sf 25x the calibrated value the block scales + clip at 448 and the codes at 6 (no NaN) -- measured, not certified. A block + whose scale is 0 dequantizes to zero whatever its codes hold, and the test's + compare does not read them. Preconditions the caller owns (a wrong preparation of each operand -- the weight's rows, its scales, the activation's scales, alpha -- runs without @@ -236,6 +222,12 @@ class CuteDslNvfp4GatherGroupedGemmActFusion(OpWrapper): global_sf of 2 elements; alpha not fp32; weight_scale typed float8_e4m3fn; input typed uint8. tile_size 64: ValueError. + Measured at bring-up and not re-driven by the test: the CUDA-graph capture + with moe_sort, the determinism and write statements, the wrong preparations' + distance from the gate, the short-alpha and short-weight_scale refusals and + the loud rejections. The test re-drives the cells, the tile / tactic + invariance and the `input` contiguity refusal. + Uncertified: Relu2 and SiTU; a SwiGLU clamp (swiglu_limit_scalar >= 0); expert_counts (DeepEP count-native metadata); expert parallelism (the kernel reads neither num_experts nor local_expert_offset, so EP changes only the sort diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_grouped_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_grouped_gemm.py index 47359870c230..3d7512323e9c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_grouped_gemm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_grouped_gemm.py @@ -30,17 +30,13 @@ def _ckpt( `mma_n` / `cluster_n` select the autotuner tactic: MMA tile `(tile_size, mma_n)`, cluster `(tile_size // 128, cluster_n)`; (128, 1) is also the fallback when the tuner has no entry. `dist` names how the test draws the - expert ids, as in the moe_sort entry's test: `router` = each row's top-10 of - uniform scores; `hot` = every row on the same 10 experts; `tile_edges` = three - experts hold exactly tile-1, tile and tile+1 rows; `round_robin` = row i - takes experts (10 i + j) mod 512. `gsf` is the FC1 `global_sf` that quantized - this op's input, against the SwiGLU output's magnitude: `calibrated` = the - output maximum the test calibrates to maps to the e4m3 maximum, as the - checkpoint's per-layer `1 / down_proj.input_scale` does; `saturating` = 25x - above that, so most input block scales clip at 448 and the codes at 6; - `subnormal` = 2^-12 of it, so input block scales fall into e4m3 subnormals - and zero. `zero_tokens` makes that many leading tokens' rows all zero, as - FC1 returns them for an all-zero token. + expert ids: `router` = each row's top-10 of uniform scores. `gsf` is the FC1 + `global_sf` that quantized this op's input, against the SwiGLU output's + magnitude: `calibrated` = the output maximum the test calibrates to maps to + the e4m3 maximum, as the checkpoint's per-layer `1 / down_proj.input_scale` + does; `subnormal` = 2^-12 of it, so input block scales fall into e4m3 + subnormals and zero. `zero_tokens` makes that many leading tokens' rows all + zero, as FC1 returns them for an all-zero token. """ return dict( tokens=tokens, @@ -127,23 +123,10 @@ class CuteDslNvfp4GroupedGemm(OpWrapper): why="tile 128: decode, 1 request without MTP; the fallback tactic", spec=_ckpt(1, 128), ), - Cell( - why="tile 128: cluster N 2 tactic, 32 requests at MTP verify width 4; largest " - "captured T", - spec=_ckpt(128, 128, cluster_n=2), - ), Cell( why="tile 128: MMA N 256 and cluster N 2 tactic, 8192-row max_num_tokens prefill", spec=_ckpt(8192, 128, mma_n=256, cluster_n=2), ), - Cell( - why="tile 128: expert rows tile-1, tile, tile+1; the tile edge", - spec=_ckpt(300, 128, dist="tile_edges"), - ), - Cell( - why="tile 128: FC1 global_sf 25x calibrated; input scales clip at 448", - spec=_ckpt(2048, 128, gsf="saturating"), - ), Cell( why="tile 128: tiny FC1 global_sf and 7 all-zero token rows; subnormal and zero " "input scales", @@ -173,8 +156,8 @@ class CuteDslNvfp4GroupedGemm(OpWrapper): 256 is certified through the per-pair invariance below. Numerics: each written element equals the reference's `out` bit for bit on - every cell, on the real FC1 output at 8192 tokens, and on the checkpoint's own - down_proj experts. Operands whose block scales span a wider range can round a + every cell (and, at bring-up, on the real FC1 output at 8192 tokens and on the + checkpoint's own down_proj experts). Operands whose block scales span a wider range can round a partial sum; the test's compare admits that at bf16 ties only. Tactics. The op asks the AutoTuner for an MMA tile of (tile_size, 128) or @@ -182,8 +165,8 @@ class CuteDslNvfp4GroupedGemm(OpWrapper): 128, 2), and with no tuning entry falls back to ((tile_size, 128), (tile_size // 128, 1)); serving warmup tunes per row bucket and may pick any of the four. Every pair's row is bitwise identical across tile 128 / 256 x - the four tactics x the fallback (the test checks it at 128 and 8192 rows), so - cells driven at one tactic certify the others. + the four tactics x the fallback (the test checks it at 8192 rows), so cells + driven at one tactic certify the others. Determinism. A (token, expert) pair's output row depends on that pair's input row and its expert alone: it is bitwise unchanged by moving it to another row @@ -228,6 +211,12 @@ class CuteDslNvfp4GroupedGemm(OpWrapper): num_non_exiting_tiles of 2 elements; tile_idx_to_group_idx int64 or one tile short. tile_size 64: ValueError. + Measured at bring-up and not re-driven by the test: the CUDA-graph capture + with moe_sort and FC1, the determinism and read / write statements, the wrong + preparations' distance from the gate, and the loud rejections. The test + re-drives the cells, the tile / tactic invariance and the `input` contiguity + refusal. + Uncertified: output dtypes other than bf16 (the runner builds none); expert parallelism (the kernel reads neither local_expert_offset nor num_experts, so EP changes only the sort and the weight slice); other hidden / intermediate diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_permuted.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_permuted.py index 9c381e0b0b65..f6efc5755b5a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_permuted.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_permuted.py @@ -58,8 +58,10 @@ class Fp4BlockScaleMoeRenormPermuted(OpWrapper): "routing kernel", spec=_ckpt(1), ), - Cell(why="16 rows: decode_first's MAX_ROWS; the cluster routing kernel", spec=_ckpt(16)), - Cell(why="128 rows: the con=128 decode batch", spec=_ckpt(128)), + Cell( + why="128 rows: the con=128 decode batch; the cluster routing kernel", + spec=_ckpt(128), + ), Cell(why="coarse 3-level logits: ties decide most slots", spec=_ckpt(64, dist="coarse")), ) @@ -81,7 +83,7 @@ class Fp4BlockScaleMoeRenormPermuted(OpWrapper): not a multiple of 256). Hard limits - T certified 1..128 (decode); the decode forward folds at + T certified at 1, 64 and 128 (decode); the decode forward folds at every site that takes the combine num_experts 512, top_k 10 """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_routed.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_routed.py index d9d50991098f..d4a76c8155c4 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_routed.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_routed.py @@ -30,7 +30,7 @@ CERTIFIED_PAIRS = frozenset({(512, 10)}) -def _ckpt(tokens: int, dist: str = "randn", logits_dtype=torch.bfloat16) -> dict: +def _ckpt(tokens: int, dist: str = "randn") -> dict: """qwen4_exp's routed geometry: 512 experts, top-10, H 2560 / I 640.""" return dict( tokens=tokens, @@ -38,7 +38,7 @@ def _ckpt(tokens: int, dist: str = "randn", logits_dtype=torch.bfloat16) -> dict intermediate=640, num_experts=512, top_k=10, - logits_dtype=logits_dtype, + logits_dtype=torch.bfloat16, dist=dist, ) @@ -61,26 +61,23 @@ class Fp4BlockScaleMoeRenormRouted(OpWrapper): ARCHS = frozenset({Arch.SM_103}) + # One cell per routing kernel the call picks by row count: block, single-cluster + # and the multi-kernel histogram path. CELLS: tuple[Cell, ...] = ( Cell(why="_moe_block decode: one token; the block routing kernel", spec=_ckpt(1)), Cell(why="128 rows: the con=128 decode batch; the cluster routing kernel", spec=_ckpt(128)), - Cell(why="8192 rows: max_num_tokens of the identity config", spec=_ckpt(8192)), Cell( why="16384 rows: max_num_tokens of the tuned config; past the 8192 tuning bucket", spec=_ckpt(16384), ), - Cell(why="coarse 3-level logits: ties decide most slots", spec=_ckpt(64, dist="coarse")), - Cell( - why="fp32 logits: the runner's other accepted logits dtype", - spec=_ckpt(64, logits_dtype=torch.float32), - ), ) note = """ Routing semantics, measured against the composition it replaces (`renorm_moe_routing_op` -> pre-routed runner): bitwise-equal finalized - output at every cell, randn / coarse / large-magnitude / all-equal logits, - bf16 and fp32 logits, 1..16384 rows. Both kernels rank by a packed + output at every cell (the test re-drives this), and at bring-up on randn / + coarse / large-magnitude / all-equal logits, bf16 and fp32 logits, + 1..16384 rows. Both kernels rank by a packed (twiddled value, 65535 - index) key (ties to the lower expert index), take an fp32 `expf` softmax over the k with the same 32-lane warp reduction, and round the weight to bf16 once. @@ -94,10 +91,13 @@ class Fp4BlockScaleMoeRenormRouted(OpWrapper): Silent, refused by `is_valid`: a non-contiguous operand (raw pointers, dense layout assumed) and a (num_experts, top_k) pair outside - `CERTIFIED_PAIRS` (a tier-list miss logs and launches nothing). + `CERTIFIED_PAIRS` (a tier-list miss logs and launches nothing; the test + holds this refusal). + + Not certified: fp32 logits (the runner's other accepted logits dtype). Hard limits - T certified 1..16384 (max_num_tokens of the tuned config). + T certified at 1, 128 and 16384 (max_num_tokens of the tuned config). Autotuner buckets stop at 8192, the op's `tune_max_num_tokens` default, which this entry does not pass: a call past 8192 rows runs the 8192 bucket's tactic diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp8_block_scale_moe_runner.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp8_block_scale_moe_runner.py index 09cbecdc20f0..2bb80ff63ecb 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp8_block_scale_moe_runner.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp8_block_scale_moe_runner.py @@ -111,15 +111,6 @@ class Fp8BlockScaleMoeRunner(OpWrapper): why="20 tokens == H/128: the activation scale buffer is square", spec=dict(tokens=20, **_CKPT), ), - Cell( - why="256 tokens: ~5 rows per expert, past the small-T routing kernel", - spec=dict(tokens=256, **_CKPT), - ), - Cell(why="1024 tokens: ~20 rows per expert", spec=dict(tokens=1024, **_CKPT)), - Cell( - why="8192 tokens: max_num_tokens of the identity config; the last tuning bucket", - spec=dict(tokens=8192, **_CKPT), - ), Cell( why="16384 rows: max_num_tokens of the tuned config; past the last tuning bucket", spec=dict(tokens=16384, **_CKPT), @@ -129,7 +120,8 @@ class Fp8BlockScaleMoeRunner(OpWrapper): note = """ Metadata consumed: none. Two process-global caches sit behind the call -- the autotuner's profiling cache and the C++ runner cache -- and a warm - tuned tactic was measured bitwise equal to the cold fallback. + tuned tactic is bitwise equal to the cold fallback (the test holds this at + 256 tokens). Preconditions. `hidden_states` is `[T, H]` e4m3 and `hidden_states_scale` `[H/128, T]` fp32 (see the class docstring); weights are the checkpoint's @@ -157,8 +149,9 @@ class Fp8BlockScaleMoeRunner(OpWrapper): `use_dp`. `gemm1_clamp_limit` is live and one scalar for every expert, applied after the FC1 requantization; None is bitwise equal to omitting it. - Silent on the raw op, refused by `is_valid`: any non-contiguous tensor (read - as dense, 55-115 ulp RMS off); `act_type != 0` (the Python runner forces + Silent on the raw op, refused by `is_valid` (the test holds the strided + `hidden_states` refusal): any non-contiguous tensor (read as dense, 55-115 + ulp RMS off); `act_type != 0` (the Python runner forces SwiGlu, so 1 and 2 are bitwise identical to 0); `routing_bias` alongside `topk_ids` (ignored; a +-1e3 bias leaves the output bitwise unchanged). `num_fused_shared_experts > 0` is rejected loudly on the pre-routed path, @@ -168,7 +161,7 @@ class Fp8BlockScaleMoeRunner(OpWrapper): states): an upstream defect, eager is unaffected. Hard limits - T certified 1..16384 (max_num_tokens of the tuned config); + T certified at 1, 20 and 16384 (max_num_tokens of the tuned config); chunk above it, not probed. Autotuner buckets stop at `tune_max_num_tokens` (default 8192, which the target leaves): past it a call runs the 8192 bucket's tactic; a @@ -176,9 +169,13 @@ class Fp8BlockScaleMoeRunner(OpWrapper): top_k <= 32 (raises past it); certified 10, the checkpoint's E, H, I certified at the checkpoint's 512, 2560, 640 - Not certified: the in-kernel routing path (`routing_logits` given), - multi-rank execution (the windows ran in one process), geometries other - than the checkpoint's, sm_100. + Not certified (the statements above were measured at bring-up and are not + re-driven by the test): expert windows (`local_expert_offset` / + `local_num_experts` other than the full stack), repeated or out-of-range + ids, `gemm1_clamp_limit`, the inert arguments, `output=`, the registered + fake, the op's loud rejections. Nor: the in-kernel routing path + (`routing_logits` given), multi-rank execution, geometries other than the + checkpoint's, sm_100. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fused_sigmoid_gate_mul_add.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fused_sigmoid_gate_mul_add.py index 473c7d9d2475..367a6139f8a8 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fused_sigmoid_gate_mul_add.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fused_sigmoid_gate_mul_add.py @@ -59,7 +59,6 @@ class FusedSigmoidGateMulAdd(OpWrapper): why="_moe_block / MTPLayer._moe decode: one token, in place; 8 warps", spec=_ckpt(1), ), - Cell(why="17 tokens: an odd multi-row count", spec=_ckpt(17)), Cell(why="2048 tokens: prefill-shaped", spec=_ckpt(2048)), Cell( why="gate as a strided column of a wider tensor, as the trunk's fused shared-expert " @@ -73,13 +72,13 @@ class FusedSigmoidGateMulAdd(OpWrapper): The Triton kernel JIT-compiles on first call. Preconditions, all enforced loudly by the callable (AssertionError before - any launch): `final_hidden_states` 2-D; `shared_expert_output` and `output` - match it in shape and dtype; `gate_logits.numel() == num_tokens`; the - destination has `stride(-1) == 1` (its row stride is free). A - column-strided `final_hidden_states` or `shared_expert_output` is densified - inside the call (an extra copy kernel), as is a non-contiguous gate. Host - operands raise ValueError from Triton. `num_tokens == 0` returns the - destination without launching. `output` may alias `shared_expert_output`. + any launch; the test holds the wrong-length gate): `final_hidden_states` + 2-D; `shared_expert_output` and `output` match it in shape and dtype; + `gate_logits.numel() == num_tokens`; the destination has `stride(-1) == 1` + (its row stride is free). A column-strided `final_hidden_states` or + `shared_expert_output` is densified inside the call (an extra copy kernel), + as is a non-contiguous gate. `num_tokens == 0` returns the destination + without launching (certified). Numerics: the default `assert_close` band holds for bf16, the dtype the target calls it with; fp16 and fp32 run the same kernel and are not certified. @@ -93,9 +92,10 @@ class FusedSigmoidGateMulAdd(OpWrapper): num_tokens row offset widened to int64 in the kernel; largest run 8192 Not certified: hidden sizes other than 2560 (the column block and warp count - step with it), an `output` buffer other than `shared_expert_output`, row-padded - operands, CUDA-graph capture, torch.compile, operands on two devices, float8 - activations. + step with it), gate layouts other than `[T, 1]` and the strided column, an + `output` buffer (including one aliasing `shared_expert_output`), row-padded + operands, host operands, CUDA-graph capture, torch.compile, operands on two + devices, float8 activations. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_sort.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_sort.py index ecfece5bfb9c..a22c0614c1d0 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_sort.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_sort.py @@ -15,11 +15,9 @@ def _ckpt(tokens: int, tile_tokens_dim: int, dist: str = "router") -> dict: """The checkpoint's trunk MoE: 512 experts, top-10, all local (tp1), fp32 scales. `dist` names how the test draws the expert ids: `router` = each row's top-10 of - uniform scores over the 512 experts; `hot` = every row routes to the same 10 - experts; `tile_edges` = three experts hold exactly tile-1, tile and tile+1 rows, - the rest `router`; `round_robin` = row i takes experts (10 i + j) mod 512; - `strided` = row i takes (h_i + 51 j) mod 512 with h_i uniform. All of them put 10 - distinct experts on every row, as the router does. + uniform scores over the 512 experts; `tile_edges` = three experts hold exactly + tile-1, tile and tile+1 rows, the rest `router`. Both put 10 distinct experts on + every row, as the router does. """ return dict( tokens=tokens, @@ -71,33 +69,21 @@ class MoeSort(OpWrapper): ARCHS = frozenset({Arch.SM_103}) + # One cell per kernel path the target reaches at its tile, 128: block (T <= 4), + # cluster (T 5..8192) and cooperative (T >= 8193). CELLS: tuple[Cell, ...] = ( Cell( why="tile 128: decode, 1 request without MTP; block kernel", spec=_ckpt(1, 128), ), Cell( - why="tile 128: 128 rows, 32 requests at MTP verify width 4; largest captured T, " - "cluster kernel", - spec=_ckpt(128, 128), - ), - Cell( - why="tile 128: expert counts tile-1, tile, tile+1; the mn_limit clamp", + why="tile 128: expert counts tile-1, tile, tile+1; the mn_limit clamp; cluster kernel", spec=_ckpt(300, 128, "tile_edges"), ), Cell( - why="tile 128: 8192 rows all to the same 10 experts; full multi-tile runs at the " - "last single-cluster T", - spec=_ckpt(8192, 128, "hot"), - ), - Cell( - why="tile 128: 8193 rows; first cooperative-kernel T", + why="tile 128: 8193 rows; first cooperative-kernel T (prefill)", spec=_ckpt(8193, 128), ), - Cell( - why="tile 256: expert counts tile-1, tile, tile+1; the other tile sm_103 admits", - spec=_ckpt(300, 256, "tile_edges"), - ), ) note = """ @@ -155,10 +141,12 @@ class MoeSort(OpWrapper): <= 0 (the host divides by it). Certified only at the checkpoint's values: num_experts 512, top_k 10, all - experts local (offset 0), tile 128 (the target's) and 256, fp32 scales, and - T up to 8193 (block, cluster and the first cooperative T). Expert - parallelism, other tiles, other expert counts and the histogram path are - uncertified. + experts local (offset 0), tile 128 (the target's), fp32 scales, and T 1, + 300 and 8193 (block, cluster and the first cooperative T); CUDA-graph replay + on new ids at T = 128; the contiguity guard. The order, unwritten-region, + data-hazard and loud-rejection statements above were measured at bring-up + and are not re-driven by the test. Expert parallelism, tile 256 and other + tiles, other expert counts and the histogram path are uncertified. Hard limits (sm_103) T * top_k boundary 2^31 pairs (int32 product) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_unpermute_inplace.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_unpermute_inplace.py index 3fddd53c8ded..e82eaa5dfdc4 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_unpermute_inplace.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_unpermute_inplace.py @@ -12,35 +12,23 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -def _ckpt( - tokens: int, - tile_tokens_dim: int = 128, - dist: str = "router", - values: str = "normal", - non_local_from: int | None = None, -) -> dict: +def _ckpt(tokens: int, values: str = "normal") -> dict: """The checkpoint's routed-expert combine: top-10 of 512 experts, hidden 2560, bf16 FC2 rows, fp32 routing weights. The test lays `permuted_input` out exactly as the CuteDSL FC2 receives it: it - sorts the ids with `moe_sort` at `tile_tokens_dim` and writes each (token, - slot) row at the permuted index the sort gave it. `dist` names how the ids are - drawn, as in the moe_sort entry's test: `router` = each row's top-10 of uniform - scores; `hot` = every row on the same 10 experts; `round_robin` = row i takes - experts (10 i + j) mod 512. `values` is how the FC2 rows are drawn: `normal` = - standard normal; `wide` = standard normal times 2^u, u uniform in [-30, 30] - per element, so one sum adds terms up to 2^60 apart in magnitude. - `non_local_from` marks every slot whose expert is at or above it `-1`, as - moe_sort does for a pair on another rank's expert, and also token 0's ten - slots. The scales are a softmax over each row's ten slots, as the - renormalizing router emits them. + draws each row's top-10 of uniform scores over the 512 experts, sorts the ids + with `moe_sort` at `tile_tokens_dim` and writes each (token, slot) row at the + permuted index the sort gave it. `values` is how the FC2 rows are drawn: + `normal` = standard normal; `wide` = standard normal times 2^u, u uniform in + [-30, 30] per element, so one sum adds terms up to 2^60 apart in magnitude. + The scales are a softmax over each row's ten slots, as the renormalizing + router emits them. """ return dict( tokens=tokens, - tile_tokens_dim=tile_tokens_dim, - dist=dist, + tile_tokens_dim=128, values=values, - non_local_from=non_local_from, num_experts=512, top_k=10, hidden_size=2560, @@ -106,24 +94,11 @@ class MoeUnpermuteInplace(OpWrapper): why="decode, 1 request without MTP; rows from the sort's block kernel", spec=_ckpt(1), ), - Cell( - why="32 requests at MTP verify width 4; largest captured T, rows placed by the " - "sort's atomics", - spec=_ckpt(128), - ), Cell(why="8192 rows, max_num_tokens prefill", spec=_ckpt(8192)), Cell( why="FC2 rows spanning 2^-30..2^30; terms 2^60 apart in one sum", spec=_ckpt(8192, values="wide"), ), - Cell( - why="510 experts with one row each; most permuted rows are padding", - spec=_ckpt(51, dist="round_robin"), - ), - Cell( - why="slots marked -1 (another rank's expert) skipped; a row of all -1 is 0", - spec=_ckpt(300, non_local_from=256), - ), ) note = """ @@ -151,7 +126,8 @@ class MoeUnpermuteInplace(OpWrapper): this op. What it writes: every element of `output`, whatever it held, and none of its inputs. - Determinism (bitwise): repeated calls; a second moe_sort of the same ids, which + Determinism (bitwise, measured at bring-up; the test re-drives only graph + replay against eager): repeated calls; a second moe_sort of the same ids, which places an expert's rows in another order; a random permutation of all permuted rows; a prefix of a larger batch, with or without a fresh sort of the prefix alone; a shuffle of the tokens. The FC2 GEMM's fused finalize @@ -170,13 +146,13 @@ class MoeUnpermuteInplace(OpWrapper): from `data_ptr()` as dense row-major); an `output` whose dtype is not `permuted_input`'s (an fp32 output receives the bf16 bit patterns in the first half of its bytes); 2^31 or more tokens (nothing is written and nothing - raises). + raises). The test holds the `permuted_input` refusal. Data hazards `is_valid` cannot see (values): a permuted index at or past `permuted_input.size(0)` is read without a bounds check; moe_sort never emits one. - Loud: int64 indices; hidden size not a multiple of 8; fp32 rows, or fp16 scales + Loud (measured, not re-driven): int64 indices; hidden size not a multiple of 8; fp32 rows, or fp16 scales with bf16 rows; fewer than `T * top_k` permuted rows; a scales or index shape that disagrees with `output`. @@ -184,8 +160,9 @@ class MoeUnpermuteInplace(OpWrapper): rows; a CPU `output`. Certified only at the checkpoint's values: top_k 10, hidden 2560, bf16 rows and - output, fp32 scales, the sort's layout at tile 128, every slot local - except in the -1 cell. + output, fp32 scales, the sort's layout at tile 128, every slot local; T 1 and + 8192 eager, 128 under CUDA-graph replay. Slots marked -1 (another rank's + expert) are not certified. Hard limits (sm_103) tokens one thread block per token, the count narrowed to int32: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/renorm_moe_routing_op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/renorm_moe_routing_op.py index 4944323b2283..c97fd51ca4e9 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/renorm_moe_routing_op.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/renorm_moe_routing_op.py @@ -48,15 +48,11 @@ class RenormMoeRoutingOp(OpWrapper): ARCHS = frozenset({Arch.SM_103}) CELLS: tuple[Cell, ...] = ( - Cell(why="_moe_block / MTPLayer._moe decode: one token, bf16 weights", spec=_ckpt(1)), + Cell(why="MTPLayer._moe decode: one token, bf16 weights", spec=_ckpt(1)), Cell( - why="4096 tokens: thousands of rows carry an exact bf16 tie in the top-10", - spec=_ckpt(4096), - ), - Cell(why="16384 rows: max_num_tokens of the tuned config", spec=_ckpt(16384)), - Cell( - why="fp32 weights at the checkpoint router, as the CuteDSL experts take them", - spec=_ckpt(512, torch.float32), + why="_moe_block at 16384 rows (max_num_tokens of the tuned config): fp32 weights, " + "as the CuteDSL experts take them; thousands of rows carry an exact bf16 tie", + spec=_ckpt(16384, torch.float32), ), Cell( why="coarse logits (3 values): ties decide most of the ranking", @@ -74,15 +70,12 @@ class RenormMoeRoutingOp(OpWrapper): Preconditions, enforced loudly by the op: `router_logits` 2-D CUDA bf16 / fp16 / fp32; `output_dtype` fp32, bf16 or None (fp16 raises even for fp16 logits); `topk <= 16`; `num_experts <= 512` -- this checkpoint sits on - that ceiling. `topk == 0` and `num_tokens == 0` return empties. + that ceiling. Silent, refused by `is_valid`: a non-contiguous `router_logits` (read as a - dense `[T, E]` buffer from `data_ptr()`) and `topk > num_experts` (the - surplus slots come back as expert 0 at weight 0 -- duplicated ids). - - Non-finite logits poison the whole row to NaN (`+inf`, `NaN`, an all - `-inf` row); a `-inf` logit among finite ones is a soft mask, at weight 0 - once `topk` exceeds the finite experts. + dense `[T, E]` buffer from `data_ptr()`; the test holds this refusal) and + `topk > num_experts` (the surplus slots come back as expert 0 at weight 0 + -- duplicated ids). `torch.ops.trtllm.default_moe_routing_op` has the identical signature and softmaxes before the top-k: same ids, un-renormalized weights (rows sum to @@ -90,15 +83,21 @@ class RenormMoeRoutingOp(OpWrapper): 512 / 16) is not bit-compatible: it uses `torch.topk` and casts before the softmax. - Bit-exact against `reference` over 2040 configurations (in dtype x out - dtype x num_experts {1..512} x topk {0..16} x tokens {0..8192}); the cells - hold the target's router (512 / top-10, bf16 logits, bf16 and fp32 weights) - and the test's sweep re-drives the boundaries of the rest of that space. + Bit-exact against `reference` over 2040 configurations at bring-up (in + dtype x out dtype x num_experts {1..512} x topk {0..16} x tokens + {0..8192}); the cells hold the target's router (512 / top-10, bf16 logits, + bf16 and fp32 weights), and the test holds the weights to HF's + `norm_topk_prob` formulation at bf16 logits. Hard limits num_experts <= 512 (raises past it) topk <= 16 (raises past it); > num_experts silent, guarded - T certified 0..16384 (max_num_tokens of the tuned config) + T certified 1..16384 (max_num_tokens of the tuned config) + + Not certified: other num_experts / topk / dtype combinations, `topk == 0` + and `num_tokens == 0` (both return empties), non-finite logits (`+inf`, + NaN and an all `-inf` row poison the row to NaN; a `-inf` among finite + logits is a soft mask). """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/flashinfer_gemma_rmsnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/flashinfer_gemma_rmsnorm.py index 4109a96da06f..4ef212a17f9b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/flashinfer_gemma_rmsnorm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/flashinfer_gemma_rmsnorm.py @@ -32,10 +32,6 @@ class FlashinferGemmaRmsnorm(OpWrapper): # 10240-wide residual, each a column slice at row stride 10240 -- that, not # a contiguous row, is what the strided-kernel cells below certify. CELLS: tuple[Cell, ...] = ( - Cell( - why="decode: 1 token, stream 3 of the 4x2560 residual (column 7680)", - spec=dict(shape=(1, 2560), col_offset=7680, **_STREAM), - ), Cell( why="decode batch 33: stream 1, row stride 10240 > width 2560", spec=dict(shape=(33, 2560), col_offset=2560, **_STREAM), @@ -48,23 +44,21 @@ class FlashinferGemmaRmsnorm(OpWrapper): why="MTP pre_e_norm: the contiguous [T, 2560] embedding", spec=dict(shape=(4, 2560), dtype=torch.bfloat16, eps=1e-6), ), - Cell( - why="hidden 5120, 8192 tokens: the largest row count measured", - spec=dict(shape=(8192, 5120), dtype=torch.bfloat16, eps=1e-6), - ), ) note = """ Metadata consumed: none. Registered only when flashinfer is importable. Certified path: flashinfer's CuTe DSL `rmsnorm_cute` kernel for 2-D bf16 - `x` (`weight_bias=1.0`), contiguous and at the residual's 10240 row stride; - a strided row selects a strided kernel variant. The 3-D per-head kernel - (`qk_rmsnorm_cute`) and fp16/fp32 operands are not certified. The CUDA-JIT fallback (`FLASHINFER_USE_CUDA_NORM=1`, + `x` (`weight_bias=1.0`) at hidden 2560: a stream slice at the residual's + 10240 row stride (33 and 2048 rows; a strided row selects a strided kernel + variant) and a contiguous `[4, 2560]` (MTP `pre_e_norm`). Not certified: a + single strided row, other widths, other row counts, the 3-D per-head kernel + (`qk_rmsnorm_cute`) and fp16/fp32 operands. The CUDA-JIT fallback (`FLASHINFER_USE_CUDA_NORM=1`, or a DSL that cannot target the device) is a different implementation and is not certified. PDL (`TRTLLM_ENABLE_PDL`, default on) changes scheduling, not results. - Loud, so not guarded: 1-D and 4-D `x`, `x.stride(-1) != 1`, a weight of + Loud (measured), so not guarded: 1-D and 4-D `x`, `x.stride(-1) != 1`, a weight of the wrong length, dtype or device, and a CPU `x` raise ValueError; float64/int8/uint8 raise KeyError from the DSL dtype table. float8_e4m3fn and float8_e5m2 are in that table and return a plausible fp8 tensor -- @@ -75,8 +69,9 @@ class FlashinferGemmaRmsnorm(OpWrapper): pointers, so odd row strides and misaligned slices were not measured. Numerics: default tolerance holds for bf16, 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. + largest deviation measured is one bf16 ulp (0.03125) at 2048 and 8192 rows + of 5120. The `+ 1` is measured: against a plain-`w` reference the output is + ~4 away, so the cells tell the two apart. Hard limits T * H contiguous inputs above 2^31 - 1 elements route to the strided diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/layer_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/layer_norm.py index b8fe44d648e1..35201417be8a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/layer_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/layer_norm.py @@ -64,25 +64,27 @@ class LayerNorm(OpWrapper): other `x`, and non-contiguous `weight`/`bias`, are silently `.contiguous()`d first (one extra copy). `weight`/`bias` dtype is free. - Certified: contiguous bf16 `[rows, 1152]` `x` with a bf16 weight and bias, - as the vision tower calls it. Not certified: strided or 3-D `x`, - `bias=None`, fp16/fp32 activations, weight/bias in another dtype than `x`. + Certified: contiguous bf16 `[rows, 1152]` `x` at 1, 1024 and 32768 rows + with a bf16 weight and bias, as the vision tower calls it. Not certified: + strided or 3-D `x`, `bias=None`, fp16/fp32 activations, weight/bias in + another dtype than `x`. Loud failures (left to the callable): weight or bias length != N raises a bare AssertionError with no message; `next_pow2(N) * x.element_size() > 64KB` raises "This layer norm doesn't support feature dim >= 64KB" (N = 65536 in bf16); CPU tensors fail before or at the Triton launch. - Accuracy: at the checkpoint's affine (weights to 17.5, biases to -25.6) the - bf16 result is the correctly rounded fp64 value on >= 99.9989% of elements; - where it differs, by <= 2^-7 of the summed term magnitude. A constant row - returns `bias` bitwise. `eps` sits inside the sqrt. + Accuracy, measured beyond the cells' default gate and not certified: at + the checkpoint's affine (weights to 17.5, biases to -25.6) the bf16 result + is the correctly rounded fp64 value on >= 99.9989% of elements; where it + differs, by <= 2^-7 of the summed term magnitude. A constant row returns + `bias` bitwise. `eps` sits inside the sqrt. Accepted but not certified: an fp8-e4m3 `x` returns fp8-e4m3 (output is `empty_like(x)`), within one e4m3 ulp of the exact value; fp64 `x`. - Forward only: the autograd Function has no backward. Runs unchanged under - `inference_mode`; CUDA-graph capturable after a warm-up call compiles the - Triton kernel. + Forward only: the autograd Function has no backward. Measured, not + certified: one kernel launch, runs unchanged under `inference_mode`, + CUDA-graph capturable after a warm-up call compiles the Triton kernel. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/rms_norm_gated_token_major.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/rms_norm_gated_token_major.py index a15d7599886e..fb79cf780dfd 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/rms_norm_gated_token_major.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/rms_norm_gated_token_major.py @@ -18,8 +18,6 @@ #: The checkpoint's GDN epilogue: 48 value heads of 128, `output_gate_type: sigmoid`. _CKPT = dict(heads=48, width=128, dtype=torch.bfloat16, eps=1e-6, gate_activation="sigmoid") -#: The same geometry under the op's default silu gate. -_SILU = dict(_CKPT, gate_activation="silu") #: Prefill-sized launches (at least this many rows, i.e. 683 tokens of 48 heads) run #: the multi-row kernel at `_PREFILL_ROWS` rows a program: at 16384 tokens 94 vs 143 #: us on GB300 (traffic roofline ~81 us), where 4 rows a program leave 196608 small @@ -80,18 +78,18 @@ class RmsNormGatedTokenMajor(OpWrapper): ARCHS = frozenset({Arch.SM_103}) CELLS: tuple[Cell, ...] = ( - *( - Cell( - why=f"z the in_proj column view: token stride 16480, {t} tokens", - spec=dict(tokens=t, z="in_proj", **_CKPT), - ) - for t in (1, 128, 8192, 16384) + Cell( + why="decode: z the in_proj column view (token stride 16480), 128 tokens", + spec=dict(tokens=128, z="in_proj", **_CKPT), + ), + Cell( + why="prefill: z the in_proj column view, 16384 tokens at 8 rows a program", + spec=dict(tokens=16384, z="in_proj", **_CKPT), ), Cell( why="683 tokens: the first launch at 8 rows a program", spec=dict(tokens=683, z="in_proj", **_CKPT), ), - Cell(why="fp8 out, scale 0.5", spec=dict(tokens=17, fp8_scale=0.5, **_SILU)), ) note = """ @@ -104,7 +102,7 @@ class RmsNormGatedTokenMajor(OpWrapper): it off. The 8-row programs reduce a row's mean square in another order than the 4-row ones (one-ulp differences on some outputs). - Loud, so not guarded: 3-D `x` (ValueError), a `z` not covering `x` + Loud (measured), so not guarded: 3-D `x` (ValueError), a `z` not covering `x` (AssertionError "does not match x shape"; 2-D `z` ValueError), an unknown `gate_activation` (ValueError before launch), and on the fallback a non-unit `x.stride(-1)` (bare AssertionError). `is_valid` covers what the @@ -116,13 +114,14 @@ class RmsNormGatedTokenMajor(OpWrapper): output is freshly allocated and contiguous; no input is mutated. Certified: bf16 `x` with the sigmoid gate over the in_proj column view of - `z` (48 heads of N = 128) on the 4-row and 8-row kernels, plus one fp8-output - cell (silu gate, packed `z`). Not certified: the generic fallback kernel, - the silu gate without fp8 output, fp16/fp32 operands, a strided `x`. + `z` (48 heads of N = 128): 128 tokens on the 4-row kernel, 683 tokens (the + first 8-row launch) and 16384 tokens on the 8-row kernel. Not certified: + fp8 output, the silu gate, the generic fallback kernel, fp16/fp32 + operands, a strided `x`. Numerics: bf16 inside the default tolerance against the fp32 reference, with one-ulp differences on ~1e-5 of elements at prefill row counts. fp8 - output byte-identical. Nothing loosened. + output byte-identical (measured). Nothing loosened. Hard limits N * itemsize fallback holds a row in 64 KB boundary 65536 B diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/gated_ngram_injection.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/gated_ngram_injection.py index cded9afd1790..af6161baa6ee 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/gated_ngram_injection.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/gated_ngram_injection.py @@ -72,7 +72,6 @@ class GatedNgramInjection(OpWrapper): # "single_channel" -- one nonzero channel per stream so each dot is an exact # product under the test's control: dots, value_dtype, clamp_min. CELLS: tuple[Cell, ...] = ( - Cell(why="checkpoint 4x2560 bf16, decode, 1 token", spec=dict(leading=(1,))), Cell( why="checkpoint 4x2560 bf16, 33 tokens / MTP verify 8 x 4 + 1", spec=dict(leading=(33,)), diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ngram_token_history.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ngram_token_history.py index 704fc63f6b14..a6f8c97dbd5e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ngram_token_history.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ngram_token_history.py @@ -93,23 +93,6 @@ class NgramTokenHistory(OpWrapper): # "literal" | "seq_lens" (zeros + int64 cumsum of an int32 seq_lens prefix -- # `cu_seqlens_from_seq_lens`' form). `ngram_size` is the checkpoint's 3. CELLS: tuple[Cell, ...] = ( - _cell( - "decode 32 x 1 token at the 461-slot engine pool", - (1,) * 32, - _scattered(32), - (True,) * 32, - ENGINE_SLOTS, - **_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=8: 32 decodes x 8 rows, qsl rebuilt", (8,) * 32, @@ -129,7 +112,6 @@ class NgramTokenHistory(OpWrapper): eos_density=0.05, **_CALL_SITE, ), - _cell("prefill 1024, fresh, EOS in stream", (1024,), (0,), (False,), 3, eos_density=0.05), _cell( "zero-token rows: carry kept, fresh reseeded", (2, 0, 0), @@ -173,13 +155,19 @@ class NgramTokenHistory(OpWrapper): pool. At 461 slots the result is bitwise identical to an 11-slot pool with the same slot content. + Certified at the forward's two verify-step shapes, both on the + `cu_seqlens_from_seq_lens` rebuild at the 461-slot pool (32 x 8 rows; a + 37-token fresh context with in-stream EOS beside 8 x 3 rows), and at + zero-token rows. Not certified: a plain decode (one row per sequence) or + a context-only batch -- outside a verify step the forward takes the + fused `ple_ngram_embed`, not this op. + CUDA-graph capture certified (no `repeat_interleave` on device lengths, so no data-dependent shape). Cost, not gated: ~0.30 ms for a 4096-token prefill, a 32-row decode or a mixed 2 x 1024 + 30 x 1 batch alike -- launch-bound. - Hard limits: not probed (`num_tokens` measured to 1024 here, 4096 in - timing). + Hard limits: not probed (`num_tokens` measured to 4096 in timing). """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_gate_conv_inject_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_gate_conv_inject_norm.py index 70e3f55490db..6afb43f1c077 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_gate_conv_inject_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_gate_conv_inject_norm.py @@ -347,10 +347,7 @@ class PleGateConvInjectNorm(OpWrapper): # max_rows (max(lengths)), carries (None -> the argument is None). The # checkpoint's 4 x 2560 bf16 geometry throughout. CELLS: tuple[Cell, ...] = ( - _cell("decode 1 x 1 row at the 461-slot pool", (1,), pool_slots=ENGINE_SLOTS), _cell("decode 128 x 1 row at the 461-slot pool", (1,) * 128, pool_slots=ENGINE_SLOTS), - _cell("MTP verify p=4: 32 x 4 rows", (4,) * 32, pool_slots=ENGINE_SLOTS), - _cell("rows shorter than max_rows, a zero-row sequence", (2, 0, 1), max_rows=4), # Context-shaped (prefill / mixed) steps: the packed-row conv. _cell("prefill 8192 rows, fresh", (8192,), carries=(False,), pool_slots=ENGINE_SLOTS), _cell( @@ -373,7 +370,12 @@ class PleGateConvInjectNorm(OpWrapper): Preconditions (data; not guarded): at most `max_rows` rows per sequence; distinct live slots; `query_start_loc[-1] == T`. Operands 2-D with unit column stride; weights contiguous. Certified at the checkpoint's 4 x 2560 - bf16 geometry with every slot live (the target never passes `pad_slot_id`). + bf16 geometry with every slot live (the target never passes `pad_slot_id`), + at a 128 x 1-row decode, an 8192-row fresh prefill and a mixed 2-context + + 64-decode step, all on the 461-slot pool. Not certified: more than one + row per sequence on the per-sequence path (`1 < max_rows <= + PER_SEQ_MAX_ROWS`) or sequences shorter than `max_rows` -- a verify step + takes the unfused ops, so the target's per-sequence path runs one row. Numerics: every rounding point of the composition is kept, so the only differences are fp32 reduction order (the three norms' sums of squares diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_ngram_embed.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_ngram_embed.py index 22f3182ff674..dea3bb7c4136 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_ngram_embed.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_ngram_embed.py @@ -328,22 +328,7 @@ class PleNgramEmbed(OpWrapper): table="ckpt", eos_density=0.05, ), - _cell("MTP verify p=4: 32 x 4 rows", (4,) * 32, pool_slots=ENGINE_SLOTS), - _cell( - "zero-row sequences republish carry", - (2, 0, 1, 0), - carries=(True, True, True, False), - max_rows=2, - ), # Context-shaped (prefill / mixed) steps: the packed-row form. - _cell( - "chunked prefill 16384 rows carrying, host table", - (16384,), - carries=(True,), - pool_slots=ENGINE_SLOTS, - table="host", - eos_density=0.01, - ), _cell( "mixed step: 2 contexts + 64 decode rows", (3001, 5000) + (1,) * 64, @@ -370,6 +355,14 @@ class PleNgramEmbed(OpWrapper): - `query_start_loc[-1] == num_tokens`; - token id * max multiplier < 2^63 (int64 wraps, as in `ngram_hash`). + Certified at a 128 x 1-row decode on a host-resident table and on the + full checkpoint-height device table, and a mixed step of 2 contexts (one + fresh, one carrying) + 64 decode rows on the packed-row form. Not + certified: more than one row per sequence on the per-sequence path + (`1 < max_rows <= PER_SEQ_MAX_ROWS`) or zero-row sequences -- a verify + step takes the unfused ops, so the target's per-sequence path runs one + row. + Bitwise equal to `ngram_token_history` + `ngram_hash` + embedding + `dequantize_e4m3_per_tensor`, window, rows, values and pool alike. CUDA-graph capturable: no host sync, one launch, output allocated per call. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/bf16_to_fp32.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/bf16_to_fp32.py index fae7a1dbaaea..61edc7eb2a1d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/bf16_to_fp32.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/bf16_to_fp32.py @@ -48,7 +48,6 @@ class Bf16ToFp32(OpWrapper): why="every bf16 bit pattern: NaNs, infinities, subnormals, signed zeros", spec=dict(shape=(1, 65536), source="all_codes"), ), - Cell(why="zero rows", spec=dict(shape=(0, VOCAB))), ) note = """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/dequantize_e4m3_per_tensor.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/dequantize_e4m3_per_tensor.py index dd2eb76c77a4..7f3922dea06a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/dequantize_e4m3_per_tensor.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/dequantize_e4m3_per_tensor.py @@ -42,7 +42,6 @@ class DequantizeE4m3PerTensor(OpWrapper): CELLS: tuple[Cell, ...] = ( Cell(why="_ple_block decode: [1, 16, 160] gathered rows, bf16 scale", spec=_ple(1)), - Cell(why="_ple_block at 8 tokens", spec=_ple(8)), Cell( why="_ple_block prefill at max_num_tokens: [16384, 16, 160] through the grid-stride " "loop", @@ -62,17 +61,14 @@ class DequantizeE4m3PerTensor(OpWrapper): grid with a grid-stride loop; there is no vectorized variant, so alignment is not an axis. - Preconditions, all enforced loudly by the op (`symmetric_dequantize_per_ - tensor` in thop/fp8Op.cpp): `input` float8_e4m3fn (e5m2, int8, bf16 raise), - CUDA, contiguous, rank >= 2, non-empty (zero rows or columns raise -- an - empty batch must branch first); `scales` CUDA, bf16 / fp16 / fp32 (fp64 - raises), of `input`'s rank with every dim 1. + Preconditions, enforced by the op (`symmetric_dequantize_per_tensor` in + thop/fp8Op.cpp): `input` float8_e4m3fn, CUDA, contiguous, rank >= 2, + non-empty (zero rows raise -- an empty batch must branch first; the test + holds this rejection); `scales` CUDA, bf16 / fp16 / fp32, of `input`'s rank + with every dim 1. - Degenerate scale values are accepted and computed exactly: 0 zeroes every - lane (signs kept, no NaN), a negative scale flips signs, +inf sends finite - lanes to +-inf and zero lanes to NaN, NaN makes everything NaN. Overflow is - inf, not a clamp (fp16 scale 1000 sends 448 to inf). NaN codes stay NaN. - `448 * PLE_SCALE` reads back as bf16 0.08935546875. + NaN codes stay NaN; `448 * PLE_SCALE` reads back as bf16 0.08935546875 (the + all-codes cell). CUDA-graph capture is certified: replay picks up a new `input` and a new `scales` value written in place (the scale is read from memory at replay). @@ -82,8 +78,9 @@ class DequantizeE4m3PerTensor(OpWrapper): int32 boundary 2^31 not probed Not certified: fp16 and fp32 outputs (the target's scale is bf16), ranks - other than 3, `TRTLLM_ENABLE_PDL=1`, `scales` on a different device than - `input`. + other than 2 and 3, degenerate scale values (0, negative, inf, NaN) and + overflow, the op's other loud rejections (dtype, layout, rank, device), + `TRTLLM_ENABLE_PDL=1`, `scales` on a different device than `input`. """ def __call__(self, input: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/fp8_quantize_1x128.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/fp8_quantize_1x128.py index 1472661e7d05..51e97e8e5444 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/fp8_quantize_1x128.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/fp8_quantize_1x128.py @@ -71,18 +71,10 @@ class Fp8Quantize1x128(OpWrapper): CELLS: tuple[Cell, ...] = ( Cell(why="MTPLayer._moe decode: one token at hidden 2560", spec=dict(tokens=1, **_CKPT)), - Cell( - why="5 tokens: the first count spilling into a second 4-row 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="8192 tokens: max_num_tokens of the identity config", - spec=dict(tokens=8192, **_CKPT), - ), Cell( why="16384 rows: max_num_tokens of the tuned config", spec=dict(tokens=16384, **_CKPT), @@ -95,27 +87,26 @@ class Fp8Quantize1x128(OpWrapper): sm_100 family the op returns a 1-D padded buffer instead. That branch was read, not run. - Preconditions, all enforced loudly by the op: `input` bf16 (fp16, fp32 and - fp8 raise), 2-D (flatten `[B, S, H]` yourself), contiguous, CUDA, and - `H % 16 == 0` -- the width, not the block count, so H need not be a - multiple of 128. `T == 0` is accepted and returns `[0, H]` / `[NB, 0]`. - A consumer needing `H % 128 == 0` (the block-scale MoE runner) owns that - constraint. - - Edge behaviour, measured: an all-zero block takes scale - `bf16(1e-10) / 448` and all-zero codes; `inf` poisons its whole block - (scale inf, its own code NaN, every other code 0); NaN stays in its own - element (the running max ignores it); the block max maps to exactly +-448 - and nothing overflows. `data` is a leading-row view of an allocation padded - to a multiple of 4 rows -- the padded rows are uninitialized, so a - raw-pointer reader must stop at `T`. CUDA-graph capturable; stream- - independent; the registered fake matches eager shapes. + Preconditions, enforced loudly by the op: `input` bf16, 2-D (flatten + `[B, S, H]` yourself), contiguous (the test holds the column-slice + rejection), CUDA, and `H % 16 == 0` -- the width, not the block count, so + H need not be a multiple of 128. A consumer needing `H % 128 == 0` (the + block-scale MoE runner) owns that constraint. + + `data` is a leading-row view of an allocation padded to a multiple of 4 + rows -- the padded rows are uninitialized, so a raw-pointer reader must + stop at `T`. CUDA-graph capture and replay are certified. Hard limits - T certified 0..16384 at H 2560 (max_num_tokens of the tuned config); - above that not probed + T certified at 1, 20 and 16384 at H 2560 (16384: max_num_tokens of the + tuned config); above that not probed H multiple of 16 (raises otherwise); certified at 2560, the draft layer's - width (the test also drives a 16-wide tail block) + width + + Not certified: `T == 0` (the op accepts it and returns `[0, H]` / + `[NB, 0]`); a short tail block (`H % 128 != 0`); all-zero blocks (the + `AMAX_FLOOR` scale), `inf` / NaN inputs and saturation; `use_ue8m0=True`; + the registered fake. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_fwd.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_fwd.py index fb73e8eda6bc..af83f235c2f7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_fwd.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_fwd.py @@ -23,13 +23,6 @@ ENGINE_SLOTS = 461 -def _engine_batch_slots() -> tuple[int, ...]: - """32 distinct slots over the whole engine pool, slot 0 and the last included.""" - slots = (0, ENGINE_SLOTS - 1) + tuple((i * 13 + 5) % (ENGINE_SLOTS - 2) + 1 for i in range(30)) - assert len(set(slots)) == 32 - return slots - - class CausalConv1dFwd(OpWrapper): """Depthwise causal conv1d (+ optional SiLU) of every sequence in a batch, against its carried window, publishing the new window -- one launch. @@ -58,10 +51,9 @@ class CausalConv1dFwd(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec keys: dim, width, dtype, lengths (varlen) or batch3d=(batch, seqlen), - # pool_slots (None -> no conv_states), slots (None -> cache_indices None), - # initial (None -> has_initial_state None), silu, bias, pad_slot_id, - # layout: "in_place" | "out" | "channel_last". + # spec keys: dim, lengths (varlen), pool_slots, slots (cache_indices), + # initial (has_initial_state); bf16, W 4, no bias, SiLU, pad_slot_id -1, + # channel-major in place throughout. CELLS: tuple[Cell, ...] = ( Cell( why="checkpoint prefill at the engine's 461-slot pool: slots 460, 0, 230", @@ -73,16 +65,6 @@ class CausalConv1dFwd(OpWrapper): 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="pad_slot_id -1 (the engine's sentinel): sequence and slot skipped", spec=dict( @@ -93,24 +75,10 @@ class CausalConv1dFwd(OpWrapper): initial=(True, True, True), ), ), - 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="one 2048-token sequence: the vectorised load, 128-thread block", spec=dict(dim=CKPT_DIM, lengths=(2048,), pool_slots=1, slots=(0,), initial=(True,)), ), - Cell( - why="one 2047-token sequence: the scalar load, 128-thread block", - spec=dict(dim=CKPT_DIM, lengths=(2047,), pool_slots=1, slots=(0,), initial=(True,)), - ), ) note = """ @@ -122,11 +90,12 @@ class CausalConv1dFwd(OpWrapper): channel-major output of `extract_transpose_prefill_slice`, W 4, no bias, SiLU on, the layer's whole conv pool unsliced, `pad_slot_id` -1, `out` None. The cells are that call (bf16, W 4, no bias, SiLU, channel-major in - place) at the engine's 461-slot pool, which tops every width the cache - manager was seen to allocate, and at small pools for the pad sentinel, - the sub-window sequences and both load forms of the 128-thread block - (single sequences of 2048 and 2047 tokens); the multi-sequence cells - run the 64-thread block with scalar loads. + place): a three-sequence prefill at the engine's 461-slot pool, which + tops every width the cache manager was seen to allocate (the 64-thread + block with scalar loads); one 2048-token sequence (the 128-thread block + with vectorised loads); and the pad sentinel on a small pool. The + 128-thread block's scalar-load form (seen at a single 2047-token + sequence) is not certified. Preconditions the op checks itself (all raise): dtypes in {bf16, fp16, fp32}, `conv_states`/`bias` dtype, int32 `query_start_loc`/`cache_indices`, @@ -137,8 +106,9 @@ class CausalConv1dFwd(OpWrapper): rather than a wrong number): a `cache_indices` value outside `[0, slots)`; `has_initial_state[s]` True with `conv_states=None` (reads through null); `conv_states`' shape (only its dtype is validated). Sequences shorter - than `W - 1` are legal and certified; a zero-length sequence, which the - target never schedules, is not certified. Unnamed slots are bitwise + than `W - 1` are legal; one (a single token without carried state) is + certified, one with carried state is not. A zero-length sequence, which + the target never schedules, is not certified. Unnamed slots are bitwise untouched. Only the channel-major kernel (the target's layout) is certified; the diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_triton_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_triton_update.py index 56437ad3cd92..cbcd77719624 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_triton_update.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_triton_update.py @@ -35,8 +35,6 @@ def _pinned(batch: int) -> tuple[int, ...]: """`batch` distinct slots of the engine pool, slot 0 and the last pinned in.""" - if batch == 1: - return (0,) return (0,) + tuple(i * 14 + 3 for i in range(1, batch - 1)) + (ENGINE_SLOTS - 1,) @@ -70,63 +68,38 @@ class CausalConv1dTritonUpdate(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec keys: dim, width (4), state_len (pool columns, 3), seqlen (2), batch, - # pool_slots, slots ("random" | tuple | None), dtype, bias, activation - # ("silu"), pad_slot_id, x_form ("token_major" | "in_proj" (token_major - # over the q|k|v columns of a [n * p, 16480] projection) | "contiguous" | "2d"), - # snapshots (rows of the snapshot buffer, or 0 for none), snapshot_steps - # (the buffer's step depth, default seqlen), accepted, launch_dependent_kernels. - CELLS: tuple[Cell, ...] = ( - tuple( - Cell( - why=f"verify, x over the [n*p, 16480] in_proj, batch {batch}, {seqlen} tok/row", - spec=dict( - dim=CKPT_DIM, - batch=batch, - seqlen=seqlen, - pool_slots=ENGINE_SLOTS, - slots=_pinned(batch), - snapshots=ENGINE_SNAPSHOT_ROWS, - x_form="in_proj", - ), - ) - for batch, seqlen in ((1, 2), (32, 3)) - ) - + 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, - ), - ) - # seqlen 8: max_draft_len 7, the widest verify step. - for batch, seqlen in ((17, 4), (32, 8)) - ) - + ( - Cell( - why="seqlen 5 into an 8-step snapshot buffer: steps 5..7 untouched", - spec=dict( - dim=CKPT_DIM, - batch=8, - seqlen=5, - pool_slots=ENGINE_SLOTS, - slots=_pinned(8), - snapshots=ENGINE_SNAPSHOT_ROWS, - snapshot_steps=8, - ), - ), - 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 - ), + # spec keys: dim, batch, seqlen, pool_slots, slots, snapshots (rows of the + # snapshot buffer), snapshot_steps (the buffer's step depth, default + # seqlen). Every cell is the call site: `x` the token-major q|k|v columns + # of a [n * p, 16480] projection seen as [n, C, p], bf16, W 4, no bias, + # SiLU, pad_slot_id -1, no accepted-token roll. + CELLS: tuple[Cell, ...] = tuple( + Cell( + why=f"verify, x over the [n*p, 16480] in_proj, batch 32, {seqlen} tok/row", + spec=dict( + dim=CKPT_DIM, + batch=32, + seqlen=seqlen, + pool_slots=ENGINE_SLOTS, + slots=_pinned(32), + snapshots=ENGINE_SNAPSHOT_ROWS, ), ) + # seqlen 8: max_draft_len 7, the widest verify step. + for seqlen in (3, 8) + ) + ( + Cell( + why="seqlen 5 into an 8-step snapshot buffer: steps 5..7 untouched", + spec=dict( + dim=CKPT_DIM, + batch=8, + seqlen=5, + pool_slots=ENGINE_SLOTS, + slots=_pinned(8), + snapshots=ENGINE_SNAPSHOT_ROWS, + snapshot_steps=8, + ), + ), ) note = """ @@ -139,10 +112,12 @@ class CausalConv1dTritonUpdate(OpWrapper): 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 the 461-slot pool - with a 32-row snapshot buffer, `x` as the target's q|k|v column view of - the projection and as a contiguous `[n, p, 10240]` block, seqlen 2, 3, 4, - 5 and 8 (a buffer deeper than seqlen keeps its later steps untouched), and - for pad-sentinel rows; bf16, W 4, no bias, SiLU only. + with a 32-row snapshot buffer and `x` as the target's q|k|v column view + of the projection: 32 rows at seqlen 3 and at seqlen 8 (the widest verify + step), and 8 rows at seqlen 5 into an 8-step buffer (a buffer deeper than + seqlen keeps its later steps untouched); bf16, W 4, no bias, SiLU only. + Not certified: a contiguous `x`, other seqlens, and pad-sentinel rows + (the target never sends one). Raise loudly, so not guarded: width 1 (triton CompilationError); a missing `intermediate_state_indices` with snapshots and slot indices given. @@ -152,8 +127,8 @@ class CausalConv1dTritonUpdate(OpWrapper): roll is uncertified; the target passes None). Out-of-range slot values: illegal access, not probed. - Accepted and inert: `launch_dependent_kernels` (bitwise identical either - way). Accepted but not exposed: `cache_seqlens` (ignored by the kernel -- + Accepted, uncertified: `launch_dependent_kernels` (the target passes + False). Accepted but not exposed: `cache_seqlens` (ignored by the kernel -- there is no ring mode), the eagle-tree mask trio, `validate_data`, `metadata`, the `_block_n`/`_num_warps`/`_num_stages` overrides. A `conv_state` dtype different from `x` is accepted and uncertified. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_update.py index a56a5017fc83..14413d38be2b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_update.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_update.py @@ -310,8 +310,6 @@ def _decode_conv(x, conv_state, weight, indices, pad_slot_id) -> None: 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,) @@ -338,31 +336,29 @@ class CausalConv1dUpdate(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec keys: dim, width (4), state_len (3), seqlen (1), batch, pool_slots, - # slots ("random" | tuple | None), dtype, bias, silu, cache_seqlens, - # pad_slot_id, x_form ("3d" | "2d" | "2d_row_slice" | "2d_in_proj" | - # "token_major"), - # strided_state. - CELLS: tuple[Cell, ...] = tuple( + # spec keys: dim, batch, pool_slots, slots ("random" | tuple), pool_layers + # (the pool a one-layer view of a buffer this many layers deep, default 1). + # Every cell is the call site: x the q|k|v column view of the decode rows + # of a [rows, 16480] in_proj output, bf16, W 4, a 3-column pool, no bias, + # SiLU, no ring, pad_slot_id -1. + CELLS: tuple[Cell, ...] = ( Cell( - why=f"x the q|k|v column view of the [rows, 16480] in_proj, batch {batch}", + why="x the q|k|v column view of the [rows, 16480] in_proj, batch 32", + spec=dict(dim=CKPT_DIM, batch=32, pool_slots=ENGINE_SLOTS, slots=_pinned(32)), + ), + Cell( + why="x the q|k|v column view of the [rows, 16480] in_proj, batch 128", + spec=dict(dim=CKPT_DIM, batch=128, pool_slots=256, slots="random"), + ), + Cell( + why="64 rows (the CuTe kernel's first) over a layer view, pad -1 rows untouched", spec=dict( dim=CKPT_DIM, - batch=batch, - pool_slots=pool, - slots=slots, - x_form="2d_in_proj", + batch=64, + pool_slots=ENGINE_SLOTS, + slots=tuple(-1 if i % 21 == 5 else i * 7 for i in range(64)), + pool_layers=2, ), - ) - for batch, pool, slots in ( - (1, ENGINE_SLOTS, _pinned(1)), - (32, ENGINE_SLOTS, _pinned(32)), - (128, 256, "random"), - ) - ) + ( - 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)), ), ) @@ -382,9 +378,10 @@ class CausalConv1dUpdate(OpWrapper): `_DecodeConvKernel` (CuTe DSL) instead: 8 channels a thread by 16-byte loads, the TRT-LLM kernel's instructions, bitwise its result (x and pool, pad rows untouched); 128 rows 7.46 -> 5.41 us a layer (GB300, graph - replay over 8 rotating layers). Certified at that call: 1 and 32 rows - over the 461-slot pool, 128 rows (the CuTe kernel) over a 256-slot pool, - and pad rows over a 128-slot pool. + replay over 8 rotating layers). Certified at that call: 32 rows (the + TRT-LLM kernel) over the 461-slot pool, 128 rows (the CuTe kernel) over a + 256-slot pool, and 64 rows (the CuTe kernel's first row count) with pad + rows over the 461-slot pool seen as one layer of a two-layer buffer. The op checks metadata itself, loudly: width 2..4, `state_len >= W - 1`, one dtype across x/conv_state/weight (bias matches weight), weight shape, @@ -397,7 +394,8 @@ class CausalConv1dUpdate(OpWrapper): Out-of-range slot values are an illegal access, not probed. A row whose index equals `pad_slot_id` (-1 at the call site) is skipped: - its x rows and the pool stay bitwise unchanged. + its x rows and the pool stay bitwise unchanged (certified on the CuTe + kernel only). Hard limits none probed. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/chunk_gated_delta_rule.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/chunk_gated_delta_rule.py index 4a5d3c03d82c..701fa42c82ed 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/chunk_gated_delta_rule.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/chunk_gated_delta_rule.py @@ -73,62 +73,24 @@ class ChunkGatedDeltaRule(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec: seqlens (packed with cu_seqlens); initial: "indexed" (a `slots`-row - # pool at `indices`) | "positional"; pool_dtype (fp32 | bf16, the dtype of - # the indexed pool); output: None | "buffer" | "window" (wide[:, 50:146] of - # a 200-token buffer). + # spec: seqlens (packed with cu_seqlens); the initial state is the + # `ENGINE_SLOTS`-row pool at `indices`, written in place; pool_dtype (fp32 | + # bf16); output: None | "buffer" (a fresh `output=` tensor). CELLS: tuple[Cell, ...] = ( Cell( why="fp32 pool (a verify step's context rows), 461 slots, output= buffer", - spec=dict( - seqlens=(200, 37, 91), - initial="indexed", - slots=ENGINE_SLOTS, - indices=_engine_slots(3), - output="buffer", - ), - ), - Cell( - why="fp32 pool, 461-slot pool, seqlens [1024]", - spec=dict( - seqlens=(1024,), initial="indexed", slots=ENGINE_SLOTS, indices=_engine_slots(1) - ), - ), - Cell( - why="bf16 pool, 461-slot pool, output= buffer", - spec=dict( - seqlens=(200, 37, 91), - initial="indexed", - slots=ENGINE_SLOTS, - indices=_engine_slots(3), - pool_dtype=torch.bfloat16, - output="buffer", - ), + spec=dict(seqlens=(200, 37, 91), indices=_engine_slots(3), output="buffer"), ), Cell( why="bf16 pool, 461-slot pool, seqlens [1024]", - spec=dict( - seqlens=(1024,), - initial="indexed", - slots=ENGINE_SLOTS, - indices=_engine_slots(1), - pool_dtype=torch.bfloat16, - ), + spec=dict(seqlens=(1024,), indices=_engine_slots(1), pool_dtype=torch.bfloat16), ), Cell( why="bf16 pool, mixed step: a context and four 1-token rows [200, 1, 1, 1, 1]", spec=dict( - seqlens=(200, 1, 1, 1, 1), - initial="indexed", - slots=ENGINE_SLOTS, - indices=_engine_slots(5), - pool_dtype=torch.bfloat16, + seqlens=(200, 1, 1, 1, 1), indices=_engine_slots(5), pool_dtype=torch.bfloat16 ), ), - Cell( - why="output= a token-window slice of a wider buffer", - spec=dict(seqlens=(96,), initial="positional", output="window"), - ), ) note = """ @@ -140,7 +102,9 @@ class ChunkGatedDeltaRule(OpWrapper): or bf16 as `mamba_ssm_cache_dtype` allocated it, int32 slots, int64 `cu_seqlens`; the MTP-verify context half passes the `si_p` / `qsl_long_p` prefixes, and a mixed batch's one-token decode rows ride in the same call. - Certified at 461 slots. + Certified at 461 slots: three ragged context rows over an fp32 pool into + an `output=` buffer, one 1024-token sequence over a bf16 pool, and a + mixed step (a context and four one-token rows) over a bf16 pool. Over a bf16 pool the initial state is read as fp32 and the final state is stored rounded once to bf16, once per call and sequence; the accumulation is unchanged. @@ -160,14 +124,16 @@ class ChunkGatedDeltaRule(OpWrapper): by the caller or the in-kernel flag (raw q/k with the flag off return finite values ~1e25). `head_first` is pinned False. - Sequence lengths are unconstrained: certified at 1..1024, off and on the - 64 boundary, mixed ragged. Memory: a `[B, NT, H, K, V]` per-chunk state - scratch in `q.dtype`, 24 MiB per 1024 tokens here. Deterministic. + Sequence lengths are unconstrained: certified at 1, 37, 91, 200 and 1024, + off and on the 64 boundary, mixed ragged. Memory: a `[B, NT, H, K, V]` + per-chunk state scratch in `q.dtype`, 24 MiB per 1024 tokens here. + Deterministic. Not certified: the FlashInfer sibling, fp16 (accepted), bf16 `g`/`beta`, an fp16 pool, `K = 256` (accepted, unmeasured), empty `T` or zero-length spans, - and what the target never passes: `initial_state=None`, `output_final_state`, - `use_qk_l2norm_in_kernel`, an explicit `scale`, no `cu_seqlens`. + and what the target never passes: `initial_state=None` or a positional + initial state, `output_final_state`, `use_qk_l2norm_in_kernel`, an explicit + `scale`, no `cu_seqlens`. Hard limits K state kernel asserts K <= 256 (read, not probed) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/cu_seqlens_from_seq_lens.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/cu_seqlens_from_seq_lens.py index 9fb98404938a..904246678bc0 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/cu_seqlens_from_seq_lens.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/cu_seqlens_from_seq_lens.py @@ -25,18 +25,11 @@ #: under MTP (2..8: `max_draft_len` 1..7); at ``p = 1`` the target reuses the #: recurrent metadata's vector instead. BATCHES = { - "decode_p2": (2,) * 32, "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), - "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`. @@ -68,12 +61,11 @@ class CuSeqlensFromSeqLens(OpWrapper): lengths=(8,) * 32, in_dtype=torch.int32, out_dtype=torch.int64, layout="prefix" ), ), - *( - Cell( - why=f"{name} {_name(torch.int32)}->{_name(torch.int64)}", - spec=dict(lengths=BATCHES[name], in_dtype=torch.int32, out_dtype=torch.int64), - ) - for name in ("chunked_prefill_p2", "padded_decode_p2", "max_batch") + Cell( + why="padded_decode_p2: zero-length padded rows, int32->int64", + spec=dict( + lengths=BATCHES["padded_decode_p2"], in_dtype=torch.int32, out_dtype=torch.int64 + ), ), ) @@ -93,9 +85,12 @@ class CuSeqlensFromSeqLens(OpWrapper): `dtype` has no default: `causal_conv1d_fwd` requires int32 and `chunk_gated_delta_rule` keys its autotuner on `cu_seqlens`' dtype. - CUDA-graph capture is certified (only `seq_lens.shape[0]` reaches Python; - a captured graph follows lengths rewritten in place). Three launches, about - 0.03 ms at batch 8..256 -- launch bound. + Certified bitwise at the call site's prefix view of `seq_lens_cuda` (a + mixed step at p 2, and the p 8 verify of 32 requests) and at a decode + batch padded with zero-length rows, int32 in and int64 out. CUDA-graph + capture is certified (only `seq_lens.shape[0]` reaches Python; a captured + graph follows lengths rewritten in place). Three launches, about 0.03 ms + at batch 8..256 -- launch bound. Not certified: a `seq_lens` dtype other than int32 (the guard admits any integer dtype), a non-contiguous `seq_lens` view. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/dilated_short_conv_paged.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/dilated_short_conv_paged.py index aec821713912..8ff846801270 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/dilated_short_conv_paged.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/dilated_short_conv_paged.py @@ -53,7 +53,10 @@ def _scattered(num_seqs: int) -> tuple[int, ...]: def _ckpt(why: str, lengths: tuple[int, ...], **spec) -> Cell: - return Cell(why=why, spec=dict(channels=CKPT_CHANNELS, lengths=lengths, **spec)) + return Cell( + why=why, + spec=dict(channels=CKPT_CHANNELS, lengths=lengths, pool_slots=ENGINE_SLOTS, **spec), + ) class DilatedShortConvPaged(OpWrapper): @@ -84,47 +87,20 @@ class DilatedShortConvPaged(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec keys: channels, kernel_size (4), dilation (3), lengths, pool_slots (8), - # slots (None -> arange), initial (default all True; "none" -> None), - # dtype, bias, silu, pad_slot_id, x_form ("channel_major" | "token_major"), - # index_dtype (int32), qsl_dtype (int32), no_indices. - CELLS: tuple[Cell, ...] = 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. + # spec keys: lengths, slots, initial (default all True). Every cell is the + # call site's own form: a bare .t() of the token-major [T, 10240] + # activation, kernel 4, dilation 3, bf16, no bias, SiLU, + # query_start_loc_long (int64), int32 slots, the whole 461-slot pool. + CELLS: tuple[Cell, ...] = ( + _ckpt("decode: 32 one-token rows", (1,) * 32, slots=_scattered(32)), _ckpt( - 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)), - ("2047+1+1+1", (2047, 1, 1, 1)), - ("32 one-token rows", (1,) * 32), - ) - ) + ( + "mixed: a 2047-token context + 3 one-token rows", (2047, 1, 1, 1), slots=_scattered(4) + ), _ckpt( - "MTP verify at 461 slots: a 37-token context + 8 rows of 3 drafts", + "MTP verify: a fresh 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, - ), - _ckpt( - "MTP verify p=8 at 461 slots: 32 rows of 8 tokens", - (8,) * 32, - pool_slots=ENGINE_SLOTS, - slots=_scattered(32), - x_form="token_major", - qsl_dtype=torch.int64, - ), - 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), ), ) @@ -137,9 +113,9 @@ class DilatedShortConvPaged(OpWrapper): dilation 3 (state 9), SiLU on, the whole `[slots, 10240, 9]` pool, int32 slots, `query_start_loc_long` (int64, rebuilt from `seq_lens` under MTP), carry flags covering the whole batch, `pad_slot_id` -1. Certified in that - form at the 461-slot pool (prefill mixes up to 2047 tokens, 32 one-token - decode rows, MTP verify batches of 3 and of 8 tokens per request, up to 32 - requests), plus a pad row on a 6-slot pool. + form at the 461-slot pool: 32 one-token decode rows, a 2047-token context + mixed with three one-token rows, and an MTP verify step (a fresh 37-token + context and 8 requests of 3 tokens). A padded row is not certified. Preconditions: `query_start_loc[-1] == x.shape[1]`; a short one indexes past the internal history and fires a device-side assert that kills the @@ -149,8 +125,9 @@ class DilatedShortConvPaged(OpWrapper): Only the checkpoint's kernel size 4 and dilation 3 are certified; there is no width dispatch. - CUDA-graph capture is certified: no host sync, all batch reads on device, - a captured graph stays correct when the batch tensors' values change. + CUDA-graph capture is certified (at 32 requests of 8 tokens): no host + sync, all batch reads on device, a captured graph stays correct when the + batch tensors' values change. Launch-bound on decode (8 rows: 0.461 ms eager vs 0.029 ms of conv; 0.094 ms replayed from a graph). diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/extract_transpose_prefill_slice.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/extract_transpose_prefill_slice.py index eebc6523d738..434e48ce4b55 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/extract_transpose_prefill_slice.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/extract_transpose_prefill_slice.py @@ -38,27 +38,17 @@ class ExtractTransposePrefillSlice(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # `num_rows`/`num_cols` build `src`; `row_stop`/`col_slice` turn it into the - # view the call site passes (a row slice of a split token buffer, or a - # column slice of a wider projection). + # `num_rows`/`num_cols` build the [T, 16480] projection; `row_stop` and + # `col_slice` turn it into the view the call site passes (its q|k|v + # columns, and in a mixed step the prefill rows of those). CELLS: tuple[Cell, ...] = ( Cell( - why="GDN prefill, whole 10240-wide row, 1 token", - spec=dict( - num_rows=1, - num_cols=GDN_MIXED_QKV_WIDTH, - n=1, - start_col=0, - width=GDN_MIXED_QKV_WIDTH, - dtype=torch.bfloat16, - ), - ), - Cell( - why="mixed batch: prefill half of a split [512, 10240] buffer, a row-slice view", + why="mixed batch: prefill rows of the q|k|v view of a [512, 16480] in_proj", spec=dict( num_rows=512, - num_cols=GDN_MIXED_QKV_WIDTH, + num_cols=GDN_IN_PROJ_WIDTH, row_stop=400, + col_slice=(0, GDN_MIXED_QKV_WIDTH), n=400, start_col=0, width=GDN_MIXED_QKV_WIDTH, @@ -85,6 +75,10 @@ class ExtractTransposePrefillSlice(OpWrapper): note = """ Metadata consumed: none. Stateless. + Certified bitwise in bf16 at the call site's operand, the q|k|v column + view of the `[T, 16480]` input projection: 33 and 8192 tokens, and a mixed + step's 400 prefill rows of a 512-row projection. + Preconditions: `src` 2-D on CUDA with `src.stride(1) == 1`; the row stride is free. A non-unit column stride (`src.t()`, `src[:, ::2]`) is refused loudly by the callable's own bare `assert`, before any launch. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/flashinfer_chunk_gated_delta_rule.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/flashinfer_chunk_gated_delta_rule.py index ec5d273c92a1..e6a503da7c0f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/flashinfer_chunk_gated_delta_rule.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/flashinfer_chunk_gated_delta_rule.py @@ -77,15 +77,16 @@ class FlashinferChunkGatedDeltaRule(OpWrapper): # qwen3.8-flash-next-nvfp4's geometry, in place over an fp32 pool. CELLS: tuple[Cell, ...] = tuple( Cell( - why=f"qwen3.8 prefill, 461-slot pool, seqlens {list(s)}", + why=f"qwen3.8 prefill into output=, 461-slot pool, seqlens {list(s)}", spec=dict( geometry=QWEN38_GEOMETRY, seqlens=s, slots=ENGINE_SLOTS, indices=_engine_slots(len(s)), + output="buffer", ), ) - for s in ((200, 37, 91), (8192,), (2048, 2048, 2048, 2048)) + for s in ((200, 37, 91), (8192,)) ) + ( Cell( why="qwen3.8 context chunk + 1-token rows [8, 1, 1, 1], 461-slot pool", @@ -96,16 +97,6 @@ class FlashinferChunkGatedDeltaRule(OpWrapper): indices=_engine_slots(4), ), ), - Cell( - why="qwen3.8 output= buffer [1, T, HV, V], 6-slot pool", - spec=dict( - geometry=QWEN38_GEOMETRY, - seqlens=(128, 64), - slots=6, - indices=(4, 1), - output="buffer", - ), - ), ) note = """ @@ -123,6 +114,10 @@ class FlashinferChunkGatedDeltaRule(OpWrapper): (`gdn_decode.gated_delta_rule_decode_pretranspose`) instead (both halves plus the concatenation). + Certified at the 461-slot fp32 pool: three ragged context rows and one + 8192-token row, each into an `output=` buffer, and an 8-token context + chunk with three one-token rows. + Memory per call: a packed fp32 copy of the named slots in and out (2 x 3 MiB per sequence), `exp(g)`, and the output. q/k/v/beta are not copied when already contiguous / fp32. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_gating.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_gating.py index d28ef96b8432..bc996ad6dd0d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_gating.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_gating.py @@ -52,26 +52,17 @@ class FusedGdnGating(OpWrapper): ARCHS = frozenset({Arch.SM_103}) # spec: rows; layout "contiguous" | "target_rows" (the decode rows of a - # contiguous [T, 48] `a` projection, past T - rows prefill rows); a_dtype, - # a_scale, fill (constant a). - 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), (8, 4), (32, 2)) - ) + ( + # contiguous [T, 48] `a` projection, past T - rows prefill rows); fill + # (constant a). `a` bf16, softplus beta 1.0, threshold 20 throughout. + CELLS: tuple[Cell, ...] = ( + Cell(why="target verify call: 32 seqs x 2 tokens, contiguous a", spec=dict(rows=64)), Cell( why="target mixed verify: decode rows of a [T, 48] projection", spec=dict(rows=16, layout="target_rows"), ), - Cell( - why="softplus beta 1.0, threshold 20", - spec=dict(rows=96, beta=1.0, threshold=20.0, a_scale=4.0), - ), Cell( why="a = 200: linear branch; exp(x) is evaluated, inf never leaks", - spec=dict(rows=4, fill=200.0, a_dtype=torch.float32), + spec=dict(rows=4, fill=200.0), ), ) @@ -81,7 +72,10 @@ class FusedGdnGating(OpWrapper): Call site (`_gdn_verify`): `(A_log, a_d, dt_bias, 1.0, 20.0)` with `a_d` the `[n * p, 48]` bf16 decode rows of the `a` projection (all of it on a pure verify step) and fp32 `A_log` / `dt_bias`; `g` is then reshaped to - `[n, p, 48]` for `fused_recurrent_gated_delta_rule_update`. + `[n, p, 48]` for `fused_recurrent_gated_delta_rule_update`. Certified at + that call: 32 sequences of 2 tokens, the 16 decode rows of a mixed step's + `[T, 48]` projection, and `a = 200` (the softplus's linear branch, where + `exp` overflows). Raise loudly, so not guarded: a non-2-D `a` (`ValueError` unpacking `a.shape`), CPU tensors (`ValueError: Pointer argument cannot be accessed diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_post_conv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_post_conv.py index 83128583a121..e34c3fdb77f1 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_post_conv.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_post_conv.py @@ -62,27 +62,18 @@ class FusedGdnPostConv(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec: num_prefill, num_decode; layout "contiguous" | "target_mixed" - # (decode a row slice of a [T, 10240] projection, a/b contiguous [T, 48]) | - # "in_proj" (the target's one [T, 16480] input projection: decode the - # q|k|v column view of its decode rows, a/b its column views); dt_bias_fill. + # spec: num_prefill, num_decode; dt_bias_fill. Every cell is the target's + # one [T, 16480] input projection: `decode` the q|k|v column view of its + # decode rows, `a`/`b` its column views, `prefill` the transposing + # extraction of its prefill rows. CELLS: tuple[Cell, ...] = ( *( Cell( why=f"target in_proj views, {p} prefill + {d} decode tokens", - spec=dict(num_prefill=p, num_decode=d, layout="in_proj"), + spec=dict(num_prefill=p, num_decode=d), ) for p, d in ((200, 32), (2048, 0)) ), - 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="prefill only, 1 token", spec=dict(num_prefill=1, num_decode=0)), - Cell( - why="mixed batch, 1 prefill + 1 decode tokens", - spec=dict(num_prefill=1, num_decode=1), - ), 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), @@ -100,7 +91,9 @@ class FusedGdnPostConv(OpWrapper): `extract_transpose_prefill_slice`'s contiguous output; `decode` and `a`/`b` column views (row stride 16480) of the layer's one `[T, 16480]` input-projection output. The stock mixer's verify call passes - `beta_dtype=b.dtype` (not certified here). + `beta_dtype=b.dtype` (not certified here). Certified in that form at a + mixed step of 200 prefill and 32 decode tokens, a 2048-token prefill-only + step, and a 64-token step with every head past the softplus threshold. Preconditions: `prefill.shape[0]` and `decode.shape[1]` >= qkv_dim (extra trailing features ignored; a short one is read out of bounds); `decode` @@ -123,8 +116,9 @@ class FusedGdnPostConv(OpWrapper): checkpoint `a + dt_bias` reaches. Not certified: fp16/fp32 activations, other head geometries, a bf16 - `beta_dtype`, bf16 `A_log`/`dt_bias`, a non-contiguous `prefill`, an empty - batch. + `beta_dtype`, bf16 `A_log`/`dt_bias`, a non-contiguous `prefill`, contiguous + `decode`/`a`/`b` (not views of the projection), single-token halves, an + empty batch. Hard limits none probed. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_cached_replay_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_cached_replay_update.py index 6edf1d4e38e8..4e6c81fe6fe5 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_cached_replay_update.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_cached_replay_update.py @@ -33,8 +33,6 @@ #: in-kernel commit, one all-layer commit after every layer has run): #: `CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE` in `modules/fla/cached_replay.py`. ALL_LAYER_COMMIT_MIN_BATCH = 16 -#: This checkpoint's GDN layer count at tp1 (36 of its 48 layers are linear attention). -NUM_GDN_LAYERS = 36 #: The recurrent pool's slot count of a full-size engine run. ENGINE_SLOTS = 461 #: Layers per slot at which the last of `ENGINE_SLOTS` slots starts past 2^31 elements of @@ -104,35 +102,42 @@ class FusedRecurrentGatedDeltaRuleCachedReplayUpdate(OpWrapper): # steps (consecutive verify steps, each followed by an acceptance per row), # accept ("random": 1..t per row, the default | "full": t every step, so from # a fresh slot at t=8 the second step appends at `pnat + t == 16` without - # wrapping and every second step after wraps | "one": 1 every step), + # wrapping and every second step after wraps), # state_dtype (the checkpoint pool's dtype: fp32 | bf16, as # `mamba_ssm_cache_dtype` allocates it; fp32 by default), # layers (each with its own activations, gates # and history; one V2-style pool interleaving the layers per slot), slots # (the pool's slot count), start ("warm": random pending histories in both # halves, a wrap and an append forced at step 0 | "fresh": pnat 0 and - # zeroed history, as the manager leaves a slot after prefill), pad (rows - # given slot -1), edge_slots (rows 0 and n-1 on the pool's last and first - # slot), output (None | "prefix" of a longer buffer | "suffix": its rows - # after 5 context tokens), packed (False: the q/k/v views over the packed - # rows without packed_qkv, as a mixed batch passes them), pdl - # (launch_with_pdl, default False, the form the target calls: PDL is a constexpr of the - # kernel, so True is a separate compile), ab ("split": separate [n, t, 48] - # projections | "strided": the column halves of one [n*t, 96] projection), - # dup (rows given the SAME slot, as CUDA-graph padding repeats one dummy slot; - # only the other rows and slots are held to the reference). + # zeroed history, as the manager leaves a slot after prefill), edge_slots + # (rows 0 and n-1 on the pool's last and first slot), dup (rows given the + # SAME slot, as CUDA-graph padding repeats one dummy slot; only the other + # rows and slots are held to the reference). Every cell passes `packed_qkv`, + # separate raw `a`/`b` projections, no `output` and `launch_with_pdl=False`. + # + # Four cells, one over the three-cell budget: the forward reaches two commit + # paths (in-kernel below 16 requests, all-layer from 16) at two checkpoint + # dtypes, and two edges no main-shape cell can carry -- CUDA-graph padding + # rows sharing one slot, and slot offsets past 2^31 elements at the engine's + # pool. CELLS: tuple[Cell, ...] = ( - # In-kernel commit (below 16 requests), one cell per token tile BT = next_pow2(T). - Cell(why="small n=1 T=2: one sequence, MTP-1", spec=_cell_spec("small", 1, 2, steps=12)), Cell( - why="small n=8 T=3 dup: rows 5..7 share one padding slot", - spec=_cell_spec("small", 8, 3, dup=(5, 6, 7), steps=6), + why="all_layer n=32 T=8 full accept: commit every second step", + spec=_cell_spec("all_layer", 32, 8, start="fresh", accept="full", steps=8), + ), + _bf16_cell( + "small n=4 T=8 full accept: pnat+T==16 appends, then wraps", + "small", + 4, + 8, + start="fresh", + accept="full", + steps=8, ), Cell( - why="small n=4 T=8 full accept: pnat+T==16 appends, then wraps", - spec=_cell_spec("small", 4, 8, start="fresh", accept="full", steps=8), + why="small n=8 T=3 dup: rows 5..7 share one padding slot", + spec=_cell_spec("small", 8, 3, dup=(5, 6, 7), steps=6), ), - # Partitioned call plus the all-layer commit (16 requests and up). Cell( why="all_layer n=16 T=3 L=8 S=461: slot 460's offset is past 2^31 elements", spec=_cell_spec( @@ -145,29 +150,6 @@ class FusedRecurrentGatedDeltaRuleCachedReplayUpdate(OpWrapper): steps=3, ), ), - Cell( - why="all_layer n=32 T=8 full accept: commit every second step", - spec=_cell_spec("all_layer", 32, 8, start="fresh", accept="full", steps=8), - ), - # ---- a bf16 checkpoint pool (`mamba_ssm_cache_dtype=bfloat16`) ---- - _bf16_cell( - "small n=4 T=8 full accept: pnat+T==16 appends, then wraps", - "small", - 4, - 8, - start="fresh", - accept="full", - steps=8, - ), - _bf16_cell( - "all_layer n=32 T=3 L=36 S=40: every GDN layer, one commit", - "all_layer", - 32, - 3, - layers=NUM_GDN_LAYERS, - slots=40, - steps=6, - ), ) note = """ @@ -227,17 +209,21 @@ class FusedRecurrentGatedDeltaRuleCachedReplayUpdate(OpWrapper): (`main_num_stages=3`, a 5-stage commit), which no cell drives. `FLA_USE_FAST_OPS=1` at import switches the kernel's `exp` to the approximate one; the cells run with it unset. - Bit for bit the same call: the in-kernel commit and the partitioned call plus - the all-layer commit; `packed_qkv` and the separate q/k/v (whose contents are - never read when `packed_qkv` is given); `launch_with_pdl` on and off; a second - run; CUDA-graph replay on new data (the test captures N = 1 at T = 2 and N = 32 at T = 8 - over an fp32 checkpoint, N = 16 at T = 4 over a bf16 one). - - T = 5..8 share one compiled token tile (`BT = next_pow2(T) = 8`); the cells drive - each tile (T = 2, 3 and 8) and both checkpoint dtypes, not every pairing of the two. At - T = 8 and history 16 a slot whose every draft is accepted appends rows 8..15 at - `pnat + T == 16` without wrapping and wraps on the next step, so it commits every second - step; the `accept="full"` cells drive that boundary (in-kernel and all-layer commit). + CUDA-graph replay on new data is bit for bit the eager call (the test captures + N = 32 at T = 8 over an fp32 checkpoint, with the all-layer commit). The two commit + paths, `packed_qkv` against separate q/k/v, and `launch_with_pdl` on and off are not + certified against each other; each form the target calls is certified against the + reference by the cells. + + Certified: the all-layer commit over an fp32 checkpoint at N = 32, T = 8 and at + N = 16, T = 3 with 8 interleaved layers over the engine's 461-slot pool (the last + slot's offset past 2^31 elements); the in-kernel commit over a bf16 checkpoint at + N = 4, T = 8 and over an fp32 one at N = 8, T = 3 with three CUDA-graph padding rows + on one shared slot. T = 5..8 share one compiled token tile (`BT = next_pow2(T) = 8`); + the cells drive the BT = 4 and BT = 8 tiles, not BT = 2 (T = 2). At T = 8 and history + 16 a slot whose every draft is accepted appends rows 8..15 at `pnat + T == 16` without + wrapping and wraps on the next step, so it commits every second step; the + `accept="full"` cells drive that boundary (in-kernel and all-layer commit). Loud refusals: work items without `use_all_layer_commit` (ValueError) and the reverse (ValueError); one of `replay_work_items` / `n_writes` alone @@ -258,8 +244,8 @@ class FusedRecurrentGatedDeltaRuleCachedReplayUpdate(OpWrapper): Not certified: pre-gated `g`/`beta` (`A_log=None`); `use_qk_l2norm_in_kernel= False`; an `ssm_states` other than fp32 or bf16, or a bf16 one at N >= 128; HV/H other - than 3, K or V other than 128; T outside 2..8 (T = 9 raises: the kernel's bare - `assert T <= 8`), at either `ssm_states` dtype; `history_size` other than 16; `old_u` / + than 3, K or V other than 128; T = 2, or T outside 2..8 (T = 9 raises: the kernel's + bare `assert T <= 8`), at either `ssm_states` dtype; `history_size` other than 16; `old_u` / `old_k` other than bf16, `old_G` other than fp32; a `scale`; any launch-knob override (`block_v`, `num_warps`, `use_l2_*`, `main_num_stages`); an output dtype other than bf16; int64 `state_indices`; N above 32. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_update.py index 6e57e396efcf..072f874f6d4e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_update.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_update.py @@ -65,13 +65,11 @@ class FusedRecurrentGatedDeltaRuleUpdate(OpWrapper): ARCHS = frozenset({Arch.SM_103}) # spec: mode ("verify": disable_state_update + buffer with cache_steps = p - # | "write_back" | "both"), n sequences x p tokens (or lengths for the - # varlen form), pool_slots (None: no pool), cache_rows, slots (None: - # arange(n)), pool_dtype / cache_dtype, beta ("fp32" | "bf16": sigmoid of - # a bf16 projection, as the target passes it | "channel"), qkv ("views" + # | "write_back"), n sequences x p tokens, pool_slots, cache_rows, slots + # (None: arange(n)), pool_dtype / cache_dtype, beta ("fp32" | "bf16": + # sigmoid of a bf16 projection, as the target passes it), qkv ("views" # over a packed [T, 10240] row block | "contiguous", as the target passes - # them), output (None | "prefix" of a longer buffer | "fp32"), - # index_dtype, cu_dtype, pool_view ("layer_major"). + # them). CELLS: tuple[Cell, ...] = tuple( _verify( n, @@ -84,11 +82,11 @@ class FusedRecurrentGatedDeltaRuleUpdate(OpWrapper): pool_dtype=dtype, cache_dtype=dtype, ) - for dtype, label, shapes in ( - (torch.float32, "", ((1, 2), (32, 2), (8, 4))), - (torch.bfloat16, "bf16 state: ", ((1, 2), (32, 2), (8, 3))), + for dtype, label, n, p in ( + (torch.float32, "", 32, 2), + (torch.float32, "", 8, 4), + (torch.bfloat16, "bf16 state: ", 32, 2), ) - for n, p in shapes ) note = """ @@ -129,15 +127,17 @@ class FusedRecurrentGatedDeltaRuleUpdate(OpWrapper): at cache_steps 2, P 4; `cache_steps` other than the buffer's step axis misplaces every row past 0 against the buffer's own indexing. - Deterministic; the first call in a fresh process is bitwise the second. + The first call in a fresh process is bitwise the second (measured on the + write-back sink; the kernel has no autotuner). Certified: the call site's form alone -- the per-step buffer sink with the pool - untouched, at 1, 8 and 32 sequences and P = 2..4, over an fp32 or a bf16 pool and buffer. - Not certified: the write-back sink (`disable_state_update=False`), alone or with the - buffer; the varlen form (`cu_seqlens`); negative, scattered or int64 indices; no pool; - fp32 or per-channel `beta`; q/k/v as views of the packed rows; an `output` buffer; - `disable_output_calculation=True`; `K` or `V` other than 128; fp16/fp32 q/k/v; P - outside 2..4; N above 32. + untouched, at 32 sequences of P = 2 and 8 of P = 4 over an fp32 pool and buffer, and + 32 of P = 2 over a bf16 one. Not certified: the write-back sink + (`disable_state_update=False`), alone or with the buffer; the varlen form + (`cu_seqlens`); negative, scattered or int64 indices; no pool; fp32 or per-channel + `beta`; q/k/v as views of the packed rows; an `output` buffer; + `disable_output_calculation=True`; `K` or `V` other than 128; fp16/fp32 q/k/v; P = 3 + and P outside 2..4; N above 32. Hard limits N * HV gridDim.z boundary 65535 probed at N=1300 (62400): runs; diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_sigmoid_gating_delta_rule_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_sigmoid_gating_delta_rule_update.py index a3d336b7b984..19e5cfdc4d9a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_sigmoid_gating_delta_rule_update.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_sigmoid_gating_delta_rule_update.py @@ -60,31 +60,26 @@ class FusedSigmoidGatingDeltaRuleUpdate(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec: seq_lens, slots (pool rows; None for no pool), indices (None -> - # arange), pool_dtype (fp32), l2norm (True), normalize_qk (caller-normalized - # q/k when l2norm is off), output ("none" | "prefix": a length-prefix slice - # of a longer buffer, as the stock mixer passes), ab ("packed": column - # halves of a [T, 96] projection | "separate": two contiguous [T, 48] | - # "in_proj": q/k/v/a/b all column views of one [T, 16480] projection, as the - # target's one input GEMM makes them), layout ("varlen" | "layer_view": - # pool[1] of a [3, slots, ...] cache | "batched": cu_seqlens None, - # contiguous copies | "rows": cu_seqlens None, one token per row, q/k/v - # [N, 1, heads, 128] views of the packed [N, 10240] block whose token - # stride is the row stride, and a/b [N, 1, 48] -- the target's bf16 call). - CELLS: tuple[Cell, ...] = tuple( + # spec: seq_lens, slots (pool rows), indices (None -> arange), pool_dtype + # (fp32), ab ("separate": two contiguous [T, 48] | "in_proj": q/k/v/a/b all + # column views of one [T, 16480] projection, as the target's one input GEMM + # makes them), layout ("varlen": [1, T, ...] with int64 cu_seqlens | + # "rows": cu_seqlens None, one token per row, q/k/v [N, 1, heads, 128] + # views of the packed [N, 10240] block whose token stride is the row + # stride, and a/b [N, 1, 48] -- the target's bf16 call). In-kernel q/k + # norm on, no output buffer. + CELLS: tuple[Cell, ...] = ( Cell( - why=f"target bf16 decode call, rows form N={n}, 461-slot pool", + why="target bf16 decode call, rows form N=128, 461-slot pool", spec=dict( - seq_lens=(1,) * n, + seq_lens=(1,) * 128, slots=ENGINE_SLOTS, - indices=_scattered(n), + indices=_scattered(128), ab="separate", layout="rows", pool_dtype=torch.bfloat16, ), - ) - for n in (1, 32, 128) - ) + ( + ), Cell( why="bf16 rows form, pad rows (-1) between slot 0 and 460", spec=dict( @@ -110,21 +105,19 @@ class FusedSigmoidGatingDeltaRuleUpdate(OpWrapper): `cu_seqlens` None -- q/k/v `[T, 1, heads, 128]` split from the whole `[T, 10240]` block reshaped to `[T, 1, 80, 128]`, so the token stride is the row stride (`is_valid` refuses other strides), `a`/`b` `[T, 1, 48]`. The - cells drive it at 1, 32 and 128 rows of a 461-slot pool, and with CUDA-graph - pad rows (-1). One more cell certifies the packed form the GDN block test + cells drive it at 128 rows of a 461-slot pool, and with CUDA-graph pad rows + (-1) at 4 rows. One more cell certifies the packed form the GDN block test uses as its fp32-pool reference: `(a_log, a, dt_bias, 1.0, 20.0, q, k, v, b, ssm_states, si_all, None, True, qsl_long, None)`, q/k/v `[1, T, heads, 128]` and `a`/`b` column views of one `[T, 16480]` input-projection output - (row stride 16480), int64 `cu_seqlens`, at 128 rows. The rows form is - bitwise the packed Triton call at either pool dtype and runs inside a CUDA - graph. + (row stride 16480), int64 `cu_seqlens`, at 128 rows. Branch: the callable takes a FlashInfer CuTe-DSL kernel (`gdn_decode_bf16_state.gated_delta_rule`) when the pool is bf16, `cu_seqlens` is given and `T == N` (head dims 128, supported arch); everything else runs the vendored Triton kernel, and only the Triton - branch is certified. The target's rows form and the reference's fp32 packed form both - stay on it, which the test proves with the FlashInfer arm armed to raise. The bf16 call takes + branch is certified. The target's rows form (no `cu_seqlens`) and the reference's + packed form over an fp32 pool both stay on it. The bf16 call takes the rows form because FlashInfer is not batch-invariant: at 48 value heads it runs one kernel up to N = 10 and another from N = 11, and one fixed row's single-step output and state differ between the two, so a @@ -156,9 +149,10 @@ class FusedSigmoidGatingDeltaRuleUpdate(OpWrapper): Not certified: the FlashInfer branch, K or V other than 128, fp16/fp32 operands, an fp16 pool, `softplus_beta`/`softplus_threshold` other than - 1.0/20.0, `B > 1` with `cu_seqlens`; the rows form above 32 rows; the packed - form over a bf16 pool, with more than one token per sequence, with an - `output` buffer, with `l2norm` off, without a pool or on a layer view of one. + 1.0/20.0, `B > 1` with `cu_seqlens`; the rows form over an fp32 pool or above + 128 rows; CUDA-graph replay against the eager call; the packed form over a + bf16 pool, with more than one token per sequence, with an `output` buffer, + with `l2norm` off, without a pool or on a layer view of one. Hard limits none probed. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_gated_norm_pdl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_gated_norm_pdl.py index 02d928a6a24c..e9179d141741 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_gated_norm_pdl.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_gated_norm_pdl.py @@ -386,17 +386,15 @@ class GdnDecodeGatedNormPdl(OpWrapper): ARCHS = frozenset({Arch.SM_103}) + # spec: rows; pool (default a contiguous 461-slot pool | "coalesced": the + # cache manager's pool, slot stride 36 layers' states, slots 100 and 101); + # pad (row 1 given slot -1). CELLS: tuple[Cell, ...] = ( Cell(why="qwen3.8 decode, 1 row (con 1), 461-slot pool", spec=dict(rows=1)), Cell(why="qwen3.8 decode, 4 rows (MAX_ROWS), scattered slots", spec=dict(rows=4)), - Cell(why="qwen3.8 decode, 2 rows, one pad slot (-1)", spec=dict(rows=2, pad=True)), Cell( - why="qwen3.8 decode, 1 row, the pool a layer view of a 3-layer cache", - spec=dict(rows=1, pool="layer_view"), - ), - Cell( - why="2 rows, the cache manager's coalesced pool: slot stride 36 layers, slot 100", - spec=dict(rows=2, pool="coalesced"), + why="2 rows, coalesced pool (slot stride 36 layers, slot 100), row 1 a pad slot (-1)", + spec=dict(rows=2, pool="coalesced", pad=True), ), ) @@ -416,6 +414,11 @@ class GdnDecodeGatedNormPdl(OpWrapper): registers. Numerics: fp32 inside; `o` rounded to bf16 before the norm, the output rounded once. Inside torch's default band for bf16 / fp32 of the fp64 composition. Deterministic. Graph-capturable after the first call per geometry (compiled then). + + Certified: 1 and 4 rows over a contiguous 461-slot pool, and 2 rows over the cache + manager's coalesced pool (slot stride 36 layers' states, slot 100 ~2.8G elements past + the base) with the second row a pad slot (-1). Under graph replay the op reads the + projection after `griddepcontrol.wait` (1 row). Not certified: 3 rows, `pdl=False`. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_rows_gated_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_rows_gated_norm.py index 9aed4c29e063..6ec2bd7904fb 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_rows_gated_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_rows_gated_norm.py @@ -456,17 +456,15 @@ class GdnDecodeRowsGatedNorm(OpWrapper): ARCHS = frozenset({Arch.SM_103}) + # spec: rows; pool (default a contiguous 461-slot pool | "coalesced": the + # cache manager's pool, slot stride 36 layers' states, slots 100 and 101); + # pad (row 1 given slot -1). CELLS: tuple[Cell, ...] = ( Cell(why="qwen3.8 decode, 128 rows (con 128), 461-slot pool", spec=dict(rows=128)), Cell(why="qwen3.8 decode, 3 rows (past the conv-fused path)", spec=dict(rows=3)), - Cell(why="qwen3.8 decode, 33 rows, one pad slot (-1)", spec=dict(rows=33, pad=True)), Cell( - why="qwen3.8 decode, 16 rows, the pool a layer view of a 3-layer cache", - spec=dict(rows=16, pool="layer_view"), - ), - Cell( - why="2 rows, the cache manager's coalesced pool: slot stride 36 layers, slot 100", - spec=dict(rows=2, pool="coalesced"), + why="2 rows, coalesced pool (slot stride 36 layers, slot 100), row 1 a pad slot (-1)", + spec=dict(rows=2, pool="coalesced", pad=True), ), ) @@ -494,6 +492,12 @@ class GdnDecodeRowsGatedNorm(OpWrapper): the transposed reduction's tree; `o` rounded to bf16 before the norm, the output rounded once. Inside torch's default band for bf16 / fp32 of the composition's float64 reference, not bitwise FlashInfer's kernel. + + Certified: 128 and 3 rows (the fewest a decode step sends here once the conv-fused + path stops) over a contiguous 461-slot pool, and 2 rows over the cache manager's + coalesced pool (slot stride 36 layers' states, slot 100 ~2.8G elements past the + base) with the second row a pad slot (-1). Under graph replay the op reads the + projection after `griddepcontrol.wait` (16 rows). Not certified: `pdl=False`. """ def __call__( @@ -598,7 +602,10 @@ class GdnDecodeRows(OpWrapper): a bf16 `out` `[rows, HV * V]` with unit column stride. The pool rows and `slots` are read before `griddepcontrol.wait`: in a mixed step nothing queued ahead of the call writes the generation rows' slots (the context rows' reset and chunked recurrence touch - only theirs). Numerics: the recurrence of `GdnDecodeRowsGatedNorm`, bitwise its `o`. + only theirs). Numerics: the recurrence of `GdnDecodeRowsGatedNorm` (the same kernel). + + Certified: a mixed step's 121 generation rows, 3 rows with a pad slot (-1), and 2 rows + over the cache manager's coalesced pool, each into a row range of the step output. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_in_proj_conv1d_pdl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_in_proj_conv1d_pdl.py index 926c3a9589f8..4110469820d5 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_in_proj_conv1d_pdl.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_in_proj_conv1d_pdl.py @@ -291,8 +291,8 @@ def launch(x, mat_b, conv_state, conv_weight, conv_state_indices, conv_dim, pad_ return out -def _cell(why: str, rows: int, slots, pdl: bool = True, pool: int = 461) -> Cell: - return Cell(why=why, spec=dict(rows=rows, slots=slots, pdl=pdl, pool=pool)) +def _cell(why: str, rows: int, slots) -> Cell: + return Cell(why=why, spec=dict(rows=rows, slots=slots, pdl=True, pool=461)) class GdnInProjConv1dPdl(OpWrapper): @@ -323,8 +323,6 @@ class GdnInProjConv1dPdl(OpWrapper): _cell("decode 1 row, slot 37 of the 461-slot pool", 1, (37,)), _cell("decode 2 rows, two slots", 2, (5, 211)), _cell("decode 2 rows, one a pad row (-1)", 2, (17, -1)), - _cell("decode 1 row, PDL off", 1, (99,), pdl=False), - _cell("decode 1 row, a 129-slot pool", 1, (128,), pool=129), ) note = """ @@ -350,10 +348,13 @@ class GdnInProjConv1dPdl(OpWrapper): fp32 over the bf16 window, taps and projection, rounded once. Only the GEMV's summation order differs from `low_m_bf16_mm` + `causal_conv1d_update`. + Certified with PDL on, over the 461-slot pool: 1 and 2 rows, and 2 rows with a pad + row; two consecutive steps under CUDA-graph replay (1 row). Not certified: `pdl=False`. + Hard limits M rows per launch boundary MAX_ROWS=2 driven at 1, 2 (cells) K a TILING key 2560 the checkpoint's hidden size - pool any slot count and strides 461 and 129 driven + pool any slot count and strides 461 driven (64 under graph replay) """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_prefill_conv_post.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_prefill_conv_post.py index 670e126997de..a24cbbf3865a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_prefill_conv_post.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_prefill_conv_post.py @@ -775,7 +775,7 @@ class GdnPrefillConvPost(OpWrapper): # spec: lengths; slots (pool rows used, None -> _engine_slots); pool_slots; # initial; layout "in_proj" (x and a/b column views of one [T, 16480] # buffer, the target's call) | "rows" (x/a/b the first rows of a longer - # projection, as a mixed step's context half); pad_slot_id; dt_bias_fill. + # projection, as a mixed step's context half), extra_rows. CELLS: tuple[Cell, ...] = ( Cell( why="target call, engine 461-slot pool: 37+1+128 tokens at slots 460, 0, 230", @@ -786,12 +786,6 @@ class GdnPrefillConvPost(OpWrapper): initial=(True, False, True), ), ), - Cell( - why="max_num_tokens 16384 over 3 contexts (the con=128 mixed step)", - spec=dict( - lengths=(5440, 5440, 5504), pool_slots=ENGINE_SLOTS, initial=(True, False, True) - ), - ), Cell( why="mixed step's context half: first 8065 rows of a 8192-row projection", spec=dict( @@ -810,22 +804,6 @@ class GdnPrefillConvPost(OpWrapper): initial=tuple(i % 3 != 0 for i in range(128)), ), ), - Cell( - why="pad_slot_id -1: that context neither reads nor writes the pool", - spec=dict( - lengths=(16, 8, 16), pool_slots=4, slots=(0, -1, 2), initial=(True, True, True) - ), - ), - Cell( - why="every head above softplus_threshold: the linear branch", - spec=dict( - lengths=(64,), - pool_slots=2, - slots=(1,), - initial=(True,), - dt_bias_fill=SOFTPLUS_THRESHOLD + 8.0, - ), - ), ) note = """ @@ -868,14 +846,19 @@ class GdnPrefillConvPost(OpWrapper): `1 + e <= 2^126`; below acc = -87.3 it returns -0 where those return a value under 1e-36. q/k norm: the 128-channel fp32 sum of squares as each thread's pairs plus a 32-lane butterfly, `sqrt.approx.ftz`, `div.full`. - Against the unfused path (entry test, at the target-call cell): - q/k/v within one bf16 ulp on 1e-6..1e-4 of elements, g/beta within the - fp32 band, conv pool bitwise. + Against the unfused path (measured at the target-call cell): q/k/v within + one bf16 ulp on 1e-6..1e-4 of elements, g/beta within the fp32 band, conv + pool bitwise. + + Certified over the 461-slot pool, seen as one layer of a two-layer buffer: + three contexts of 37, 1 and 128 tokens; a mixed step's 8065-token context + half, the first rows of an 8192-row projection; and 128 short contexts (1 to + 64 tokens, some shorter than the window). A `pad_slot_id` context reads no carry and publishes nothing; its outputs are computed from a zero window (the unfused conv leaves those columns unconvolved instead -- nothing consumes either). The engine gives every - context row a real slot. + context row a real slot; a pad context is not certified. Preconditions (`is_valid`): `x` 2-D, unit channel stride, at least `qkv_dim` wide; `a`/`b` at least `T` rows (read through both strides); @@ -892,11 +875,12 @@ class GdnPrefillConvPost(OpWrapper): Not certified: fp16/fp32 activations, other head geometries, bias, a row stride of `x` other than the projection's 16480 (the conv pass compiles one - kernel per row stride), zero-token contexts. + kernel per row stride), zero-token contexts, `a + dt_bias` past the softplus + threshold on every head. Hard limits element offsets 64-bit for the rows of x and the outputs; weight and - pool strides int32 (refused from 2^31) not probed beyond 16384 tokens. + pool strides int32 (refused from 2^31) not probed beyond 8192 tokens. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/pack_gdn_decode_qkv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/pack_gdn_decode_qkv.py index 5ff86c1442bf..8fe4119dbaaa 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/pack_gdn_decode_qkv.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/pack_gdn_decode_qkv.py @@ -57,7 +57,7 @@ class PackGdnDecodeQkv(OpWrapper): 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, 4), (32, 2)) + for b, p in ((32, 2), (8, 4)) ) + (Cell(why="127 tokens, contiguous: a partial final block", spec=dict(num_tokens=127)),) note = """ @@ -66,7 +66,9 @@ class PackGdnDecodeQkv(OpWrapper): Call site (`_gdn_verify`): `mixed_qkv` the `[n * p, 10240]` token-major conv output (contiguous: the conv result carries its input's strides, so - transposing back is a view), heads `16, 128, 48, 128`, bf16. + transposing back is a view), heads `16, 128, 48, 128`, bf16. Certified bitwise + at that call for 32 sequences of 2 tokens and 8 of 4, and at 127 contiguous + rows (a partial final block). Raise loudly, so not guarded: `head_k_dim` and `head_v_dim` rounding up to different powers of two (`128 / 64`: triton `CompilationError` diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/reset_recurrent_state_rows.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/reset_recurrent_state_rows.py index 086581080989..b61b97d3bb67 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/reset_recurrent_state_rows.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/reset_recurrent_state_rows.py @@ -59,9 +59,9 @@ class ResetRecurrentStateRows(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec: slots (pool rows), indices, flags (bool), index_dtype (int32), - # rec_dtype (fp32), conv (True), layout: "plain" | "row_padding" - # (stride(0) > row) | "layer_view" (pool[1] of a [3, slots, ...] allocation). + # spec: slots (pool rows), indices (int32), flags (bool), rec_dtype (fp32), + # layout: "plain" | "layer_view" (pool[1] of a [3, slots, ...] allocation). + # Both pools are passed throughout. CELLS: tuple[Cell, ...] = ( Cell( why="engine pool 461, 32 requests, slot 460 reset", @@ -71,16 +71,6 @@ class ResetRecurrentStateRows(OpWrapper): why="bf16 recurrent pool 461, 32 requests, slot 0 reset", spec=dict(_engine_case(32, flip=False), rec_dtype=torch.bfloat16), ), - Cell( - why="negative and out-of-range slots skipped, int32", - spec=dict( - slots=6, indices=(-1, 6, 23, -99999, 2), flags=(False,) * 5, index_dtype=torch.int32 - ), - ), - 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"), @@ -91,8 +81,10 @@ class ResetRecurrentStateRows(OpWrapper): 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 4, 6 - and 461 slots, over an fp32 and a bf16 recurrent pool. + `[layers, slots, ...]` allocations, passed unsliced. Certified at 461 slots + over an fp32 and a bf16 recurrent pool, and at 4 slots over one layer's view + of a 3-layer allocation. Not certified: negative or out-of-range slots, + rows with padding between them, `conv_states=None`. Loud, from the callable itself (ValueError): an empty recurrent pool, conv and recurrent slot counts differing (the in-range test uses the recurrent diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/patch_embed_conv3d.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/patch_embed_conv3d.py index 561418cadd79..16b691b4008a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/patch_embed_conv3d.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/patch_embed_conv3d.py @@ -54,7 +54,7 @@ class PatchEmbedConv3d(OpWrapper): why=f"checkpoint geometry, {n} patches bf16", spec=dict(num_patches=n, dtype=torch.bfloat16), ) - for n in (1, 64, 4096) + for n in (64, 4096) ) + ( Cell( why="tower call ceiling: 8192 tokens x merge 2x2 = 32768 patches bf16", @@ -83,8 +83,8 @@ class PatchEmbedConv3d(OpWrapper): Not certified: dtypes other than bf16 (the fp32 TF32 behavior above is measured, not certified); other geometries, `bias=None` and non-contiguous - pixel rows; patch counts between the 4096- and 32768-row cells, and past - 32768. + pixel rows; patch counts below 64, between the cells, and past 32768 + (cuDNN picks its kernel by row count). """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/pos_embed_interpolate.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/pos_embed_interpolate.py index 5fc3075f782b..b616a48a86de 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/pos_embed_interpolate.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/pos_embed_interpolate.py @@ -28,15 +28,6 @@ HIDDEN = 1152 MERGE = 2 -#: Grids the processor produces: the smallest, the exact-fit grid, the -#: non-square serving grid and the largest one call carries (32768 patches). -CHECKPOINT_GRIDS = ( - ((1, 2, 2),), - ((1, 48, 48),), - ((1, 62, 90),), - ((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.""" @@ -67,15 +58,7 @@ class PosEmbedInterpolate(OpWrapper): 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 - ) + ( + CELLS: tuple[Cell, ...] = ( Cell( why="a 4-image batch incl. the 62x90 serving grid", spec=dict( @@ -96,6 +79,17 @@ class PosEmbedInterpolate(OpWrapper): hidden=HIDDEN, ), ), + # The largest grid one call carries (32768 patches). + Cell( + why="checkpoint table bf16, grid [[1, 128, 256]]", + spec=dict( + grid=((1, 128, 256),), + merge=MERGE, + dtype=torch.bfloat16, + rows=NUM_GRID**2, + hidden=HIDDEN, + ), + ), ) note = """ @@ -116,8 +110,11 @@ class PosEmbedInterpolate(OpWrapper): at capture. Cost is ~225 us of launch-bound index construction per image, flat in token count; the gather is the only part that scales. - Not certified: a CPU table (it runs and returns on CPU); fp16 and fp32 - tables; a `spatial_merge_size` other than the checkpoint's 2. + Certified at the checkpoint's bf16 table on a 4-image batch (4x8, a 2x6x4 + video, the 48x48 exact fit and the 62x90 serving grid), a 3-frame 6x10 + video and the 128x256 grid one call carries at most. Not certified: a CPU + table (it runs and returns on CPU); fp16 and fp32 tables; a + `spatial_merge_size` other than the checkpoint's 2. """ def __call__( diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py index 71bc78290455..1cfb3de32ee3 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py @@ -52,12 +52,6 @@ def randn(*shape): def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name certifying(op) out, gate_buffer, gate = _build(cell.spec, seed=_seed(cell.why)) - 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, - ) gate_before = gate_buffer.clone() expected = op.reference(out, gate) @@ -68,66 +62,11 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res 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[0].spec, tokens=17) - out, q_and_gate, gate = _build(spec, seed=7) - q_half = q_and_gate[:, :, : spec["head_dim"]] - - from_strided = op(out.clone(), gate) - from_compacted = op(out.clone(), gate.contiguous()) - from_wrong_half = op(out.clone(), q_half) - - assert torch.equal(from_strided, from_compacted) - torch.testing.assert_close(from_strided, op.reference(out, gate)) - assert (from_wrong_half.float() - from_strided.float()).abs().max().item() > 0.5 - - -@pytest.mark.parametrize( - "case", - [ - "output_3d", - "output_last_stride", - "gate_numel_mismatch", - "gate_token_mismatch", - "gate_2d_shape_mismatch", - "gate_1d", - "gate_4d", - "gate_last_stride", - ], -) -def test_metadata_misuse_raises(case: str) -> None: - """Every layout violation is loud, which is why the entry has no `is_valid`.""" +def test_metadata_misuse_raises() -> None: + """A gate with a strided last dim is refused loudly, as is every layout + violation; hence the entry has no `is_valid`.""" g = torch.Generator(device="cuda").manual_seed(8) - - def randn(*shape): - return torch.randn(*shape, generator=g, device="cuda").to(torch.bfloat16) - - attention_output, gate = randn(4, 128), randn(4, 2, 64) - if case == "output_3d": - attention_output = randn(4, 2, 64) - elif case == "output_last_stride": - attention_output = randn(4, 256)[:, ::2] - elif case == "gate_numel_mismatch": - gate = randn(4, 2, 32) - elif case == "gate_token_mismatch": - gate = randn(3, 2, 64) - elif case == "gate_2d_shape_mismatch": - gate = randn(4, 64) - elif case == "gate_1d": - gate = randn(512) - elif case == "gate_4d": - gate = randn(4, 2, 2, 32) - elif case == "gate_last_stride": - gate = randn(4, 256)[:, ::2] - + attention_output = torch.randn(4, 128, generator=g, device="cuda").to(torch.bfloat16) + gate = torch.randn(4, 256, generator=g, device="cuda").to(torch.bfloat16)[:, ::2] with pytest.raises(AssertionError): op(attention_output, gate) - - -def test_a_cpu_tensor_is_rejected_by_triton() -> None: - attention_output = torch.randn(4, 128, dtype=torch.bfloat16) - gate = torch.randn(4, 128, dtype=torch.bfloat16) - with pytest.raises(ValueError, match="Pointer argument cannot be accessed from Triton"): - op(attention_output, gate) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_erf.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_erf.py index 9beeec3d0c99..5255d4197ee8 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_erf.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_erf.py @@ -139,75 +139,3 @@ def test_identity_is_the_erf_form_and_not_tanh() -> None: print(f" bf16: correct {correct_b:.6e}; tanh variant {tanh_b:.6e} ({tanh_b / correct_b:.3f}x)") assert tanh_b < 4.0 * correct_b, (tanh_b, correct_b) torch.testing.assert_close(outb.to(torch.bfloat16), op.reference(xb)) - - -def test_peak_divergence_between_the_two_gelus() -> None: - """Where and how far apart the two GELUs are: the gap bounds every claim - about which one an implementation is.""" - xs = torch.linspace(-8.0, 8.0, 2_000_001, device="cuda", dtype=torch.float64) - gap = (_erf64(xs) - _tanh64(xs)).abs() - i = int(gap.argmax().item()) - print(f"\n max|erf - tanh| = {gap[i].item():.6e} at x = {xs[i].item():.4f}") - assert 4.0e-4 < gap[i].item() < 1e-3 - - -def test_extremes_and_non_finite() -> None: - """Saturation, signed-zero underflow, no overflow at finfo.max, -inf -> nan.""" - for dtype in (torch.bfloat16, torch.float16, torch.float32): - finfo = torch.finfo(dtype) - probe = torch.tensor( - [0.0, -0.0, 1e-3, -1e-3, 20.0, -20.0, 1000.0, -1000.0, finfo.max, -finfo.max], - device="cuda", - dtype=dtype, - ) - out = op(probe) - torch.testing.assert_close(out, op.reference(probe)) - assert out[4].item() == 20.0 and out[6].item() == 1000.0 - assert out[8].item() == finfo.max, "no overflow at the top of the dtype's range" - for index in (5, 7, 9): - assert out[index].item() == 0.0 and math.copysign(1.0, out[index].item()) < 0 - - nonfinite = torch.tensor( - [float("inf"), float("-inf"), float("nan")], device="cuda", dtype=dtype - ) - got = op(nonfinite) - assert got[0].item() == float("inf") - assert math.isnan(got[1].item()), "gelu_erf(-inf) is nan" - assert math.isnan(got[2].item()) - ref = _erf64(nonfinite) - assert ref[0].item() == float("inf"), "the closed form agrees: the function, not the kernel" - assert math.isnan(ref[1].item()) and math.isnan(ref[2].item()) - - -def test_rank_does_not_change_the_arithmetic() -> None: - """The one engine op computing this function rejects non-2-D input; this one does not.""" - x = _uniform((2, 37, MERGER_WIDTH), -16.0, 16.0, torch.bfloat16, 26) - assert torch.equal(op(x).reshape(-1, MERGER_WIDTH), op(x.reshape(-1, MERGER_WIDTH))) - - -def test_rows_are_independent() -> None: - """The 4-to-1 grouping in front of the merger reorders rows; the activation - must not care. Permuting rows permutes the result, bitwise.""" - x = _uniform((140, MERGER_WIDTH), -16.0, 16.0, torch.bfloat16, 27) - order = torch.randperm(140, generator=torch.Generator().manual_seed(0)).to("cuda") - assert torch.equal(op(x[order].contiguous()), op(x)[order]) - - -def test_cuda_graph_capture_and_replay() -> None: - x = _uniform((64, MERGER_WIDTH), -16.0, 16.0, torch.bfloat16, 28) - - warmup = torch.cuda.Stream() - warmup.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(warmup): - for _ in range(3): - op(x) - torch.cuda.current_stream().wait_stream(warmup) - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - captured = op(x) - - x.copy_(_uniform((64, MERGER_WIDTH), -16.0, 16.0, torch.bfloat16, 29)) - graph.replay() - torch.cuda.synchronize() - torch.testing.assert_close(captured, op.reference(x)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_tanh.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_tanh.py index c8a7fa602db3..b09168720ce5 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_tanh.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_tanh.py @@ -131,63 +131,3 @@ def test_identity_is_the_tanh_form_and_not_erf() -> None: print(f" bf16: correct {correct_b:.6e}; erf variant {erf_b:.6e} ({erf_b / correct_b:.3f}x)") assert erf_b < 4.0 * correct_b, (erf_b, correct_b) torch.testing.assert_close(outb.to(torch.bfloat16), op.reference(xb)) - - -def test_extremes_and_non_finite() -> None: - """Saturation, signed-zero underflow, no overflow at finfo.max, -inf -> nan.""" - for dtype in (torch.bfloat16, torch.float16, torch.float32): - finfo = torch.finfo(dtype) - probe = torch.tensor( - [0.0, -0.0, 1e-3, -1e-3, 20.0, -20.0, 1000.0, -1000.0, finfo.max, -finfo.max], - device="cuda", - dtype=dtype, - ) - out = op(probe) - torch.testing.assert_close(out, op.reference(probe)) - assert out[4].item() == 20.0 and out[6].item() == 1000.0 - assert out[8].item() == finfo.max, "x**3 is not evaluated in the input dtype" - for index in (5, 7, 9): - assert out[index].item() == 0.0 and math.copysign(1.0, out[index].item()) < 0 - - nonfinite = torch.tensor( - [float("inf"), float("-inf"), float("nan")], device="cuda", dtype=dtype - ) - got = op(nonfinite) - assert got[0].item() == float("inf") - assert math.isnan(got[1].item()), "gelu_tanh(-inf) is nan" - assert math.isnan(got[2].item()) - ref = _tanh64(nonfinite) - assert ref[0].item() == float("inf"), "the closed form agrees: the function, not the kernel" - assert math.isnan(ref[1].item()) and math.isnan(ref[2].item()) - - -def test_rank_does_not_change_the_arithmetic() -> None: - x = _uniform((2, 37, INTERMEDIATE_SIZE), -16.0, 16.0, torch.bfloat16, 26) - assert torch.equal(op(x).reshape(-1, INTERMEDIATE_SIZE), op(x.reshape(-1, INTERMEDIATE_SIZE))) - - -def test_rows_are_independent() -> None: - """Permuting rows permutes the result, bitwise.""" - x = _uniform((128, INTERMEDIATE_SIZE), -16.0, 16.0, torch.bfloat16, 27) - order = torch.randperm(128, generator=torch.Generator().manual_seed(0)).to("cuda") - assert torch.equal(op(x[order].contiguous()), op(x)[order]) - - -def test_cuda_graph_capture_and_replay() -> None: - x = _uniform((64, INTERMEDIATE_SIZE), -16.0, 16.0, torch.bfloat16, 28) - - warmup = torch.cuda.Stream() - warmup.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(warmup): - for _ in range(3): - op(x) - torch.cuda.current_stream().wait_stream(warmup) - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - captured = op(x) - - x.copy_(_uniform((64, INTERMEDIATE_SIZE), -16.0, 16.0, torch.bfloat16, 29)) - graph.replay() - torch.cuda.synchronize() - torch.testing.assert_close(captured, op.reference(x)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_silu_and_mul_rows.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_silu_and_mul_rows.py index 2fe01861e216..e5291a35eb15 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_silu_and_mul_rows.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_silu_and_mul_rows.py @@ -1,8 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """GPU test for the silu_and_mul_rows entry: the cells bitwise `flashinfer_silu_and_mul` on a -contiguous copy of the gate|up columns (and so inside its reference band), the extra columns -unread, graph capture.""" +contiguous copy of the gate|up columns (and so inside its reference band) with the fused +output's extra columns present, graph capture, the width guard.""" import zlib @@ -44,13 +44,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by assert torch.equal(out, flashinfer_silu_and_mul(x[:, : 2 * INTER].contiguous())) -def test_extra_columns_are_not_read() -> None: - x = _build(dict(rows=64), 7) - a = op(x, INTER) - x[:, 2 * INTER :] = float("nan") - assert torch.equal(op(x, INTER), a) - - def test_cuda_graph_capture_and_replay() -> None: x = _build(dict(rows=128), 8) warmup = torch.cuda.Stream() @@ -71,5 +64,3 @@ def test_guards() -> None: x = _build(dict(rows=4), 5) with validating(op), pytest.raises(AssertionError): op(x[:, : 2 * INTER - 2], INTER) - with validating(op), pytest.raises(AssertionError): - op(x.float(), INTER) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_canonical_selection.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_canonical_selection.py index 46b72e2c21fe..0561d585eeae 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_canonical_selection.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_canonical_selection.py @@ -26,9 +26,6 @@ COMPRESS_RATIO = 4 BLOCK_TOPK = 512 COLS = 2080 -#: What a padding slot sorts under when a raw selection is put in ascending order: past -#: every column, so padding stays last. -PAD_KEY = torch.iinfo(torch.int32).max def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: @@ -51,33 +48,19 @@ def _logits(kind: str, rows: int, cols: int, gen: torch.Generator) -> torch.Tens if kind == "relu_sparse": x = torch.relu(torch.randn(rows, cols, generator=gen, device=DEV)) * 3.0 return x * (torch.rand(rows, cols, generator=gen, device=DEV) < 0.2) - if kind == "signed_zeros": - x = _logits("relu_sparse", rows, cols, gen) - negative = torch.rand(rows, cols, generator=gen, device=DEV) < 0.5 - return torch.where((x == 0) & negative, torch.full_like(x, -0.0), x) - if kind == "pair_at_k": - x = torch.rand(rows, cols, generator=gen, device=DEV).argsort(1).float() + 1.0 - order = torch.argsort(x[:, :2048], dim=1, descending=True) - kth, nxt = order[:, BLOCK_TOPK - 1 : BLOCK_TOPK], order[:, BLOCK_TOPK : BLOCK_TOPK + 1] - x.scatter_(1, nxt, torch.gather(x, 1, kth)) - return x raise ValueError(kind) def _windows(spec: dict, rows: int, cols: int, gen: torch.Generator): kind = spec["windows"][0] starts = torch.zeros(rows, dtype=torch.int32, device=DEV) - if kind in ("chunk", "positions_random"): - if kind == "chunk": - positions = spec["windows"][1] + torch.arange(rows, device=DEV) - else: - lo, hi = spec["windows"][1:] - positions = torch.randint(lo, hi + 1, (rows,), generator=gen, device=DEV) - ends = ((positions + 1) // COMPRESS_RATIO).clamp_max(cols).to(torch.int32) + if kind == "chunk": + positions = spec["windows"][1] + torch.arange(rows, device=DEV) else: - assert kind == "ends", kind - ends = torch.tensor(spec["windows"][1], dtype=torch.int32, device=DEV) - assert starts.numel() == rows and ends.numel() == rows + assert kind == "positions_random", kind + lo, hi = spec["windows"][1:] + positions = torch.randint(lo, hi + 1, (rows,), generator=gen, device=DEV) + ends = ((positions + 1) // COMPRESS_RATIO).clamp_max(cols).to(torch.int32) return starts, ends @@ -135,13 +118,6 @@ def _build(spec: dict, rows: int, seed: int): return logits, starts, ends, indices -def _sorted(raw: torch.Tensor) -> torch.Tensor: - """A raw selection in ascending order, padding last.""" - keyed = torch.where(raw < 0, PAD_KEY, raw) - ordered = torch.sort(keyed, dim=1).values - return torch.where(ordered == PAD_KEY, -1, ordered) - - @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) @@ -154,46 +130,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 out = op(*args) assert out is args[3], "the rewrite writes `indices` and returns it" _compare(out, expected) - if not spec.get("graph"): - return - - # Captured once, replayed against new scores, new causal bounds and a new engine - # selection through the same buffers. - logits, starts, ends, indices = _build(spec, rows, seed) - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - op(logits, starts, ends, indices) - for replay in range(2): - fresh = _build(spec, rows, seed + 1000 + replay) - for static, new in zip((logits, starts, ends, indices), fresh): - static.copy_(new) - expected = op.reference(logits, starts, ends, indices) - graph.replay() - torch.cuda.synchronize() - _compare(indices, expected) - - -def test_tie_free_rows_equal_the_radix_engines_sorted_selection() -> None: - """Where the k-th value has no tie, the rewrite is the engine's own set in ascending - order, bit for bit.""" - gen = torch.Generator(device=DEV).manual_seed(4) - rows = 256 - logits = torch.randn(rows, COLS, generator=gen, device=DEV) - positions = torch.randint(0, 8192, (rows,), generator=gen, device=DEV) - ends = ((positions + 1) // COMPRESS_RATIO).clamp_max(COLS).to(torch.int32) - starts = torch.zeros(rows, dtype=torch.int32, device=DEV) - _poison(logits, starts, ends, "nan") - window = torch.where( - torch.arange(COLS, device=DEV)[None, :] < ends[:, None], logits, float("-inf") - ) - top = torch.topk(window, BLOCK_TOPK + 1, dim=1).values - long_rows = ends > BLOCK_TOPK - assert bool((top[long_rows, -2] != top[long_rows, -1]).all()), "a tie at k: reseed" - assert int(long_rows.sum()) > 100 and int((~long_rows).sum()) > 10 - raw = _radix(logits, starts, ends) - expected = _sorted(raw) - assert torch.equal(op(logits, starts, ends, raw), expected) def test_narrowing_the_table_to_a_bound_on_the_windows_changes_nothing() -> None: @@ -212,36 +148,6 @@ def test_narrowing_the_table_to_a_bound_on_the_windows_changes_nothing() -> None assert torch.equal(wide, narrow) -@pytest.mark.parametrize("kind", ["relu_sparse", "pair_at_k"]) -def test_radix_launches_differ_and_canonical_rows_do_not(kind) -> None: - """Repeated radix launches over one tie-heavy score matrix. The raw rows differ from - launch to launch -- in order, and where the tie class at the k-th score straddles the - budget, in which tied columns are kept -- and the canonical row is one and the same - row, the reference's, for every launch of every row. Zero ties as the target's relu - scores have them, and a positive pair tied exactly at the k-th value.""" - gen = torch.Generator(device=DEV).manual_seed(17) - rows, launches = 64, 20 - logits = _logits(kind, rows, COLS, gen) - starts = torch.zeros(rows, dtype=torch.int32, device=DEV) - ends = torch.full((rows,), 2048, dtype=torch.int32, device=DEV) - _poison(logits, starts, ends, "nan") - expected = op.reference(logits, starts, ends, torch.empty(rows, BLOCK_TOPK, device=DEV)) - orders = [set() for _ in range(rows)] - sets = [set() for _ in range(rows)] - canonical = [set() for _ in range(rows)] - for _ in range(launches): - raw = _radix(logits, starts, ends) - for r, (row, ordered) in enumerate(zip(raw.cpu().tolist(), _sorted(raw).cpu().tolist())): - orders[r].add(tuple(row)) - sets[r].add(tuple(ordered)) - out = op(logits, starts, ends, raw) - for r, row in enumerate(out.cpu().tolist()): - canonical[r].add(tuple(row)) - _compare(out, expected) - assert any(len(o) > 1 for o in orders), "radix returned one order on every launch" - assert all(len(c) == 1 for c in canonical) - - def test_guard_refuses_a_broadcast_row_start() -> None: """A one-element `row_starts` broadcasts over every row without an error and frames every row by row 0's start. Shown unguarded first, then refused.""" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_expand_block_indices.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_expand_block_indices.py index 0904d0565654..0c83dcf90793 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_expand_block_indices.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_expand_block_indices.py @@ -31,8 +31,6 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: DEV = "cuda" CR, TOPK = 4, 2048 -BLOCK_TOPK = TOPK // CR -FINAL_TOPK = TOPK + CR - 1 ENGINE_MAX_SEQ_LEN = 8192 @@ -63,14 +61,13 @@ def _rows(batch) -> tuple[list[int], list[int]]: positions = [4096 + 13 * i for i in range(batch[1])] elif kind == "serving": return _serving_rows(batch[1], batch[2]) - elif kind == "verify": + else: + assert kind == "verify", kind starts, per_request = batch[1], batch[2] positions = [s + i for s in starts for i in range(per_request)] lengths = [s + per_request for s in starts for _ in range(per_request)] assert max(lengths) <= ENGINE_MAX_SEQ_LEN return positions, lengths - else: - positions = list(batch[1]) return positions, [p + 1 for p in positions] @@ -95,17 +92,9 @@ def _packed_selection(positions: torch.Tensor, compress_ratio: int, token_topk: def _build(spec, seed: int): positions, lengths = _rows(spec["batch"]) cr, topk = spec["compress_ratio"], spec["token_topk"] - block_topk = topk // cr query_positions = torch.tensor(positions, dtype=spec["position_dtype"], device=DEV) sequence_lengths = torch.tensor(lengths, dtype=spec["length_dtype"], device=DEV) blocks = _packed_selection(query_positions, cr, topk, seed) - rows = blocks.shape[0] - 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 = torch.randint(0, 64, (rows + 37, block_topk), dtype=torch.int32, device=DEV) - tall[:rows] = blocks - blocks = tall[:rows] return (blocks, query_positions, sequence_lengths), dict(compress_ratio=cr, token_topk=topk) @@ -126,89 +115,10 @@ def test_certified_cells(cell, certifying): # noqa: F811 ) -def _expand(blocks, positions, lengths, *, compress_ratio=CR, token_topk=TOPK): - """One call, gated on the entry's own reference, returned on the host.""" - args = ( - blocks.to(DEV), - torch.as_tensor(positions, dtype=torch.int64).to(DEV), - torch.as_tensor(lengths, dtype=torch.int32).to(DEV), - ) - geometry = dict(compress_ratio=compress_ratio, token_topk=token_topk) - out = op(*args, **geometry) - _compare(out, op.reference(*args, **geometry)) - return out.cpu() - - -def test_worked_example_small_geometry() -> None: - """Two rows written out by hand, end to end: the reference agrees too.""" - # Row 0: groups 1 and 0 -> tokens 4..7 then 0..3; position 9 leaves {8, 9}. - # Row 1: group 0 -> 0..3; position 5 leaves {4, 5}; the rest is padding. - blocks = torch.tensor([[1, 0], [0, -1]], dtype=torch.int32) - out = _expand(blocks, [9, 5], [10, 6], compress_ratio=4, token_topk=8) - assert out.tolist() == [ - [4, 5, 6, 7, 0, 1, 2, 3, 8, 9, -1], - [0, 1, 2, 3, 4, 5, -1, -1, -1, -1, -1], - ] - - -def test_launch_shapes_agree_on_shared_rows() -> None: - """The same rows expand identically either side of the 128-row switch.""" - positions = torch.tensor([4096 + 13 * i for i in range(128)], device=DEV) - blocks = _packed_selection(positions, CR, TOPK, seed=7) - whole_row = _expand(blocks, positions, positions + 1) - column_split = _expand(blocks[:127], positions[:127], positions[:127] + 1) - assert torch.equal(whole_row[:127], column_split) - - -@pytest.mark.parametrize( - "position, expected_tail", - [(4098, 3), (4099, 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: the - widest tail, none, and the first row. Every width 0..3 is a certified cell.""" - 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] - - -@pytest.mark.parametrize("rows", (4, 256)) -def test_cuda_graph_capture_and_replay(rows) -> None: - """Capturable, and a replay recomputes from whatever the buffers now hold: 4 rows, and the - 256 rows a verify step of 32 requests carries at max_draft_len 7.""" +def test_cuda_graph_capture_and_replay() -> None: + """Capturable, and a replay recomputes from whatever the buffers now hold: the 256 rows + a verify step of 32 requests carries at max_draft_len 7.""" + rows = 256 positions = torch.tensor([4096 + 7 * i for i in range(rows)], device=DEV) lengths = (positions + 1).to(torch.int32) blocks = _packed_selection(positions, CR, TOPK, seed=31) @@ -236,74 +146,11 @@ def call(): _compare(captured, op.reference(blocks, positions, lengths, compress_ratio=CR, token_topk=TOPK)) -def test_output_is_fresh_and_inputs_are_untouched() -> None: - """Every column is written, two calls do not alias, and nothing is in place.""" - positions = torch.tensor([4096, 97, 4095], device=DEV) - lengths = (positions + 1).to(torch.int32) - blocks = _packed_selection(positions, CR, TOPK, seed=41) - before = (blocks.clone(), positions.clone(), lengths.clone()) - outputs = [op(blocks, positions, lengths, compress_ratio=CR, token_topk=TOPK) for _ in range(5)] - assert len({out.data_ptr() for out in outputs}) == len(outputs) - # The output comes from torch.empty: a skipped column would differ between - # calls or land outside the legal range. - for out in outputs[1:]: - assert torch.equal(outputs[0], out) - first = outputs[0] - assert bool(((first == -1) | ((first >= 0) & (first < int(lengths.max())))).all()) - for tensor, saved in zip((blocks, positions, lengths), before): - assert torch.equal(tensor, saved) - - -def test_op_rejects_malformed_arguments() -> None: - """Every metadata violation is the callable's own loud error: why is_valid is empty. - One case per error family.""" +def test_op_rejects_a_width_that_does_not_match_the_geometry() -> None: + """A selection buffer whose width is not `token_topk // compress_ratio` is the + callable's own loud error, not a silent mis-expansion: why is_valid is empty.""" positions = torch.tensor([4096, 4097], device=DEV) lengths = torch.tensor([4097, 4098], dtype=torch.int32, device=DEV) blocks = _packed_selection(positions, CR, TOPK, seed=17) - geometry = dict(compress_ratio=CR, token_topk=TOPK) - cases = [ - (r"must have shape \[rows, block_topk\]", (blocks.reshape(-1), positions, lengths), {}), - ("divisible by compress_ratio >= 2", (blocks, positions, lengths), dict(compress_ratio=1)), - ("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), {}), - ( - "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), {}), - ] - for message, args, override in cases: - with pytest.raises(ValueError, match=message): - op(*args, **{**geometry, **override}) - - -def test_discriminates_plausible_wrong_variants() -> None: - """The bitwise gate separates this layout from the ones a rewrite would produce.""" - positions = [4096, 97, 4095] - blocks = _packed_selection(torch.tensor(positions, device=DEV), CR, TOPK, seed=23).cpu() - out = _expand(blocks, positions, [p + 1 for p in positions]) - first_free = (blocks >= 0).sum(dim=1) * CR - - dropped_tail = out.clone() - tail_at_fixed_column = out.clone() - for row, position in enumerate(positions): - dropped_tail[row, first_free[row] :] = -1 - tail_at_fixed_column[row, first_free[row] : TOPK] = -1 - tail_start = (position + 1) // CR * CR - for offset in range(min(position + 1 - tail_start, CR - 1)): - tail_at_fixed_column[row, TOPK + offset] = tail_start + offset - assert not torch.equal(out, dropped_tail) - assert not torch.equal(out, tail_at_fixed_column) - assert not torch.equal(out, torch.sort(out, dim=1).values) + with pytest.raises(ValueError, match="width does not match"): + op(blocks, positions, lengths, compress_ratio=CR, token_topk=1024) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_mtp_reuse_topk.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_mtp_reuse_topk.py index a8746a2040e2..1f8e34122860 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_mtp_reuse_topk.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_mtp_reuse_topk.py @@ -35,7 +35,6 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: DEV = "cuda" CR = 4 -K = 512 #: The engine's `max_num_tokens`: height of the request-map and Top-K buffers. ENGINE_MAX_NUM_TOKENS = 8192 ENGINE_MAX_SEQ_LEN = 8192 @@ -57,18 +56,12 @@ def _captured_row(visible: int, block_topk: int, gen: torch.Generator) -> torch. def _requests(contexts, drafted, gen) -> list[tuple[int, int]]: """(captured_visible, current_visible) per request.""" kind = contexts[0] - if kind == "gaps": - return list(contexts[1]) if kind == "explicit": positions = list(contexts[1]) else: count = contexts[1] - low, high = {"mixed": (0, 8188), "short": (0, 2046), "full": (2047, 8188)}[kind] - fixed = { - "mixed": ANCHORS, - "short": (2, 2046, 2045, 2044, 2043, 402, 3), - "full": (2047, 2050, 8188, 2051), - }[kind] + low, high = {"mixed": (0, 8188), "full": (2047, 8188)}[kind] + fixed = {"mixed": ANCHORS, "full": (2047, 2050, 8188, 2051)}[kind] positions = list(fixed[:count]) positions += torch.randint(low, high + 1, (count - len(positions),), generator=gen).tolist() out = [] @@ -106,11 +99,8 @@ def _build(spec, seed: int, requests: int | None = None) -> dict: shared[slot, q] = _captured_row(captured, block_topk, gen) shared_vis[slot, q] = captured - order = list(range(num_requests)) - if spec["request_map"] == "permuted": - order = torch.randperm(num_requests, generator=gen).tolist() request, visible = [], [] - for q in order: + for q in range(num_requests): captured, current = pairs[q] for j in range(per_request): request.append(q) @@ -119,34 +109,21 @@ def _build(spec, seed: int, requests: int | None = None) -> dict: request = torch.tensor(request, dtype=meta_dtype).reshape(rows) visible = torch.tensor(visible, dtype=meta_dtype).reshape(rows) - if spec["layout"] == "engine": - height = max(ENGINE_MAX_NUM_TOKENS, rows + 37) - tall_out = torch.randint(-1, 3000, (height, block_topk), generator=gen).to(torch.int32) - tall_req = torch.randint(0, shared_rows, (height,), generator=gen).to(meta_dtype) - tall_vis = torch.randint(0, 3000, (height,), generator=gen).to(meta_dtype) - tall_req[:rows], tall_vis[:rows] = request, visible - buffers = dict( - shared=shared.to(DEV), - shared_vis=shared_vis.to(DEV), - request=tall_req.to(DEV), - visible=tall_vis.to(DEV), - output=tall_out.to(DEV), - ) - else: - assert spec["layout"] == "dense" - buffers = dict( - shared=shared[slot].clone().to(DEV), - shared_vis=shared_vis[slot].clone().to(DEV), - request=request.to(DEV), - visible=visible.to(DEV), - output=torch.full((rows, block_topk), -7, dtype=torch.int32, device=DEV), - ) - slot = None + height = max(ENGINE_MAX_NUM_TOKENS, rows + 37) + tall_out = torch.randint(-1, 3000, (height, block_topk), generator=gen).to(torch.int32) + tall_req = torch.randint(0, shared_rows, (height,), generator=gen).to(meta_dtype) + tall_vis = torch.randint(0, 3000, (height,), generator=gen).to(meta_dtype) + tall_req[:rows], tall_vis[:rows] = request, visible + buffers = dict( + shared=shared.to(DEV), + shared_vis=shared_vis.to(DEV), + request=tall_req.to(DEV), + visible=tall_vis.to(DEV), + output=tall_out.to(DEV), + ) kwargs = dict( - shared_indices=buffers["shared"] if slot is None else buffers["shared"][slot], - shared_visible_blocks=buffers["shared_vis"] - if slot is None - else buffers["shared_vis"][slot], + shared_indices=buffers["shared"][slot], + shared_visible_blocks=buffers["shared_vis"][slot], request_indices=buffers["request"][:rows], visible_blocks=buffers["visible"][:rows], output=buffers["output"][:rows], @@ -239,212 +216,33 @@ def _spec_of(prefix: str) -> dict: return cell.spec -def test_boundary_rows_take_the_documented_branch() -> None: - """The boundary cell's rows land where the contract says, read off the output.""" - spec = _spec_of("8 requests, one per boundary") - built = _build(spec, seed=5) - kwargs = built["kwargs"] - out = op(**kwargs).cpu() - shared = kwargs["shared_indices"].cpu() - # (captured, current) per request: (0,1) (1,1) (100,101) (511,512) (512,513) - # (2047,2048) (513,513) (512,512) - assert out[0, 0] == 0 and bool((out[0, 1:] == -1).all()), "empty row: group 0 at slot 0" - assert torch.equal(out[1], shared[1]), "no new group: a plain copy" - assert torch.equal(out[2, :100], shared[2, :100]) and out[2, 100] == 100 - assert bool((out[2, 101:] == -1).all()), "short row: appended right after the prefix" - assert torch.equal(out[3, :511], shared[3, :511]) and out[3, 511] == 511, "fills slot 511" - for row, group in ((4, 512), (5, 2047)): - assert torch.equal(out[row, :511], shared[row, 1:]) and out[row, 511] == group, ( - "full row: the new group displaces the oldest id and the rest shift to the front" - ) - assert torch.equal(out[6], shared[6]) and torch.equal(out[7], shared[7]) - - -def test_fresh_two_rows_take_the_documented_branch() -> None: - """Draft steps 4..6 (max_draft_len 7) append up to two groups; read off the output. - - A full row loses its two oldest ids and the survivors shift to the front, a 511-id row - loses its oldest id, and a short row keeps every captured id with the new groups right - after the prefix. At fresh 1 a full row gives up only the id in slot 0. - """ - spec = _spec_of("8 requests at fresh 2") - kwargs = _build(spec, seed=23)["kwargs"] - out = op(**kwargs).cpu() - shared = kwargs["shared_indices"].cpu() - captured = kwargs["shared_visible_blocks"].cpu().tolist() - current = kwargs["visible_blocks"].cpu().tolist() - # (captured, current) per request: (512, 514) (100, 102) (511, 513) (2046, 2048) (1, 3) - # (512, 513) (255, 257) (512, 514) - assert [c - a for a, c in zip(captured, current)] == [2, 2, 2, 2, 2, 1, 2, 2] - for row in (0, 3, 7): - first = captured[row] - assert torch.equal(out[row, : K - 2], shared[row, 2:]), "full row: the oldest two go" - assert out[row, K - 2 :].tolist() == [first, first + 1], "full row: the two new groups last" - assert torch.equal(out[2, : K - 2], shared[2, 1 : K - 1]), "511-id row: the oldest id goes" - assert out[2, K - 2 :].tolist() == [511, 512] - assert torch.equal(out[1, :100], shared[1, :100]) and out[1, 100:102].tolist() == [100, 101] - assert bool((out[1, 102:] == -1).all()), "short row: appended right after the prefix" - assert out[4, :3].tolist() == [0, 1, 2] and bool((out[4, 3:] == -1).all()) - assert torch.equal(out[5, : K - 1], shared[5, 1:]) and out[5, K - 1] == 512, ( - "fresh 1 on a full row: only the id in slot 0 goes" - ) - assert torch.equal(out[6, :255], shared[6, :255]) and out[6, 255:257].tolist() == [255, 256] - assert bool((out[6, 257:] == -1).all()) - - -def test_worked_example_small_geometry() -> None: - """Five rows written out by hand at K = 8, through a permuted request map.""" - shared = torch.tensor( - [ - [5, 2, 7, 0, -1, -1, -1, -1], # 4 of 8 groups, 2 more completed - [3, 1, 6, 4, 0, 2, 7, 5], # full, 1 more: displaces its first id - [-1] * 8, # nothing captured yet - [1, 0, 2, -1, -1, -1, -1, -1], # nothing new: a copy - [4, 1, 0, 3, 2, 6, 5, -1], # 7 ids + 2 new: displaces its first id too - ], - dtype=torch.int32, - device=DEV, - ) - captured = torch.tensor([8, 8, 0, 3, 7], dtype=torch.int32, device=DEV) - request = torch.tensor([1, 0, 4, 2, 3], dtype=torch.int32, device=DEV) - current = torch.tensor([9, 10, 9, 1, 3], dtype=torch.int32, device=DEV) - kwargs = dict( - shared_indices=shared, - shared_visible_blocks=captured, - request_indices=request, - visible_blocks=current, - output=torch.full((5, 8), 99, dtype=torch.int32, device=DEV), - ) - out = op(**kwargs) - _compare(out, op.reference(**kwargs)) - assert out.cpu().tolist() == [ - [1, 6, 4, 0, 2, 7, 5, 8], - [5, 2, 7, 0, 8, 9, -1, -1], - [1, 0, 3, 2, 6, 5, 7, 8], - [0, -1, -1, -1, -1, -1, -1, -1], - [1, 0, 2, -1, -1, -1, -1, -1], - ] - - -def test_unpacked_prefix_overwrites_a_valid_id() -> None: - """The append slot is the count of valid ids, so a -1 hole costs a group, silently. - - Prefix packing is a value property, so no metadata guard can check it. - """ - shared = torch.tensor([[3, -1, 1, -1, -1, -1, -1, -1]], dtype=torch.int32, device=DEV) - kwargs = dict( - shared_indices=shared, - shared_visible_blocks=torch.tensor([4], dtype=torch.int32, device=DEV), - request_indices=torch.tensor([0], dtype=torch.int32, device=DEV), - visible_blocks=torch.tensor([5], dtype=torch.int32, device=DEV), - output=torch.empty((1, 8), dtype=torch.int32, device=DEV), - ) - out = op(**kwargs) - _compare(out, op.reference(**kwargs)) - assert out.cpu().tolist() == [[3, -1, 4, -1, -1, -1, -1, -1]] - - -def _guard_case(name: str): - """Legal arguments, then the one layout `is_valid` refuses. +def _guard_case(): + """Legal arguments, then the whole `[slots, seqs]` visible-count buffer passed where + one slot's view belongs. Returns (intended, kwargs): `intended` is the reference on the legal arguments -- what a caller passing the bad layout meant to get. """ - spec = _spec_of("permuted request map") + spec = _spec_of("decode MTP") kwargs = _build({**spec, "slots": (3, 2)}, seed=13, requests=6)["kwargs"] intended = op.reference(**kwargs) - rows = kwargs["output"].shape[0] - - def strided(tensor): - wide = torch.zeros((tensor.numel() * 2,), dtype=tensor.dtype, device=DEV) - wide[::2] = tensor.reshape(-1) - return wide[::2] - - if name == "output column slice": - kwargs["output"] = torch.zeros((rows, 2 * K), dtype=torch.int32, device=DEV)[:, :K] - elif name == "shared_indices 3-D, middle extent == block_topk": - # Slot 1 meant; the 3-D buffer passes the width check and reads slot 0. - stack = torch.full((2, 8, 8), -1, dtype=torch.int32, device=DEV) - stack[1, :, :2] = torch.tensor([5, 6], dtype=torch.int32, device=DEV) - intended_kwargs = dict( - shared_indices=stack[1], - shared_visible_blocks=torch.full((8,), 7, dtype=torch.int32, device=DEV), - request_indices=torch.arange(4, dtype=torch.int32, device=DEV), - visible_blocks=torch.full((4,), 8, dtype=torch.int32, device=DEV), - output=torch.empty((4, 8), dtype=torch.int32, device=DEV), - ) - intended = op.reference(**intended_kwargs) - kwargs = {**intended_kwargs, "shared_indices": stack} - elif name == "shared_visible_blocks 2-D [slots, seqs]": - whole = torch.zeros((3, kwargs["shared_visible_blocks"].numel()), dtype=torch.int32) - whole = whole.to(DEV) - whole[2] = kwargs["shared_visible_blocks"] - kwargs["shared_visible_blocks"] = whole - else: - field = name.split()[1] - kwargs[field] = strided(kwargs[field]) + whole = torch.zeros((3, kwargs["shared_visible_blocks"].numel()), dtype=torch.int32) + whole = whole.to(DEV) + whole[2] = kwargs["shared_visible_blocks"] + kwargs["shared_visible_blocks"] = whole return intended, kwargs -#: One case per guard family: a non-contiguous output, strided row metadata, and a -#: whole multi-slot buffer passed where one slot's view belongs (indices and counts). -GUARDS = ( - ("output column slice", "output must be contiguous"), - ("shared_indices 3-D, middle extent == block_topk", "shared_indices must be one slot's"), - ("shared_visible_blocks 2-D [slots, seqs]", "shared_visible_blocks must be one slot's"), - ("strided request_indices", "request_indices must be contiguous"), -) - - -@pytest.mark.parametrize("name, message", GUARDS, ids=[g[0] for g in GUARDS]) -def test_guard_refuses_a_silently_wrong_layout(name: str, message: str) -> None: - """Unguarded the op returns a plausible wrong answer; `is_valid` refuses it.""" - intended, kwargs = _guard_case(name) +def test_guard_refuses_a_whole_multi_slot_count_buffer() -> None: + """Unguarded the op reads slot 0's counts and returns a plausible wrong answer; + `is_valid` refuses it.""" + intended, kwargs = _guard_case() out = op(**kwargs) torch.cuda.synchronize() - assert not torch.equal(out, intended), f"{name}: the op got it right after all" - _, kwargs = _guard_case(name) - with validating(op), pytest.raises(ValueError, match=message): + assert not torch.equal(out, intended), "the op got it right after all" + _, kwargs = _guard_case() + with ( + validating(op), + pytest.raises(ValueError, match="shared_visible_blocks must be one slot's"), + ): op(**kwargs) - - -def test_op_rejects_malformed_arguments() -> None: - """The callable's own loud errors, one per family: why `is_valid` does not restate them.""" - spec = {**_spec_of("8 requests, one per boundary"), "layout": "dense"} - kwargs = _build(spec, seed=17)["kwargs"] - cases = [ - ("width does not match", dict(shared_indices=kwargs["shared_indices"][:, : K // 2])), - ( - "one request and bound per query row", - dict(request_indices=kwargs["request_indices"][:3]), - ), - ("requires CUDA tensors", dict(shared_indices=kwargs["shared_indices"].cpu())), - ] - for message, override in cases: - with pytest.raises(ValueError, match=message): - op(**{**kwargs, **override}) - - -def test_discriminates_plausible_wrong_variants() -> None: - """The bitwise gate separates this contract from the ones a rewrite would produce.""" - spec = _spec_of("8 requests, one per boundary") - kwargs = _build(spec, seed=19)["kwargs"] - out = op(**kwargs).cpu() - rows = kwargs["request_indices"].long().cpu() - shared = kwargs["shared_indices"].cpu()[rows] - captured = kwargs["shared_visible_blocks"].long().cpu()[rows] - current = kwargs["visible_blocks"].long().cpu() - - never_displace = shared.clone() # appends only where there is room - always_last = shared.clone() # always writes the last `fresh` slots - displace_last = shared.clone() # a full row gives up its newest ids, not its oldest - for r in range(out.shape[0]): - fresh = int((current[r] - captured[r]).clamp(0, K)) - filled = int((shared[r] >= 0).sum()) - for j in range(fresh): - if filled + j < K: - never_displace[r, filled + j] = captured[r] + j - always_last[r, K - fresh + j] = captured[r] + j - displace_last[r, min(filled, K - fresh) + j] = captured[r] + j - for wrong in (never_displace, always_last, displace_last, shared): - assert not torch.equal(out, wrong.to(torch.int32)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_scores.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_scores.py index 62aa3a5b936f..8498ad6cded2 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_scores.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_scores.py @@ -37,9 +37,8 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: 2.9e-6 over every cell (2.4e-6 at 8191 rows x 2080 columns, scores to 11.6), inside the default rtol 1.3e-6 / atol 1e-5. Wrong variants (no relu, mean over heads, first group member as key, unscaled, - 1/head_dim scale, block table ignored) land more than 1000x outside - (asserted at 2080 columns; measured 3e5x to 4e6x). Columns the - reference marks NaN are unspecified and not read. + 1/head_dim scale, block table ignored) were measured 3e5x to 4e6x + outside it. Columns the reference marks NaN are unspecified and not read. """ assert out.dtype == ref.dtype == torch.float32, (out.dtype, ref.dtype) assert out.shape == ref.shape, (out.shape, ref.shape) @@ -52,7 +51,6 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: DECODE_ROW_LIMIT = 256 #: Index heads must divide this into more than one row for the query tile. DOT_TILE_WIDTH = 128 -CAPTURE_PAGES = 260 class _LaunchRecorder: @@ -101,7 +99,7 @@ def grid(self): def _case(rows, requests, **overrides): - """A narrow-table spec for the non-cell tests, same shape as the entry's own.""" + """A narrow-table spec for the non-cell tests, same geometry as the entry's own.""" spec = dict( layout="random", rows=rows, @@ -130,28 +128,17 @@ def _rows_and_requests(spec, gen): def pages_for(seq_len): return math.ceil(seq_len / tpb) - if layout in ("random", "arange"): + if layout == "random": rows, requests = spec["rows"], spec["requests"] req = torch.arange(rows, device=DEV) * requests // max(rows, 1) - if layout == "arange": - positions = torch.arange(rows, device=DEV) - else: - positions = torch.randint(0, pages * tpb, (rows,), generator=gen, device=DEV) + positions = torch.randint(0, pages * tpb, (rows,), generator=gen, device=DEV) return positions, req, [pages] * requests - segments = [] # (first position, row count) per request, in row order + # (first position, row count) per request, in row order if layout == "capture": segments = [(0, n) for n in spec["lengths"]] elif layout == "decode": segments = [(p, spec["rows_per_request"]) for p in spec["gen_positions"]] - elif layout == "mixed": - max_position = pages * tpb - 1 - contexts = spec["context_lens"] - segments = [ - (0 if i == 0 else (907 * i) % (max_position - n + 2), n) for i, n in enumerate(contexts) - ] - # MTP: each generation request carries two rows at consecutive positions. - segments += [(100 + (499 * g) % 8000, 2) for g in range(spec["requests"] - len(contexts))] else: raise AssertionError(f"unknown layout {layout!r}") positions = torch.cat([torch.arange(s, s + n, device=DEV) for s, n in segments]) @@ -245,127 +232,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 _compare(out, op.reference(**args)) -@pytest.mark.parametrize("context_rows,branch", [(False, "row_wise"), (True, "query_tile")]) -def test_capture_width_needs_more_than_one_column_program(context_rows, branch) -> None: - """Why the 260-page table is an axis of its own, not a bigger number. - - Both kernels walk the score columns in a grid dimension of their own. At 8 - pages (64 columns) one program covers every column; at 260 pages (2080) it - does not, so only the wide cells ever run a later column program. - """ - grids = {} - for tag, pages in (("narrow", 8), ("wide", CAPTURE_PAGES)): - args = _build(_case(300, 4, pages=pages, context_rows=context_rows), seed=8500 + pages) - with _watch_launches() as watcher: - out = op(**args) - torch.cuda.synchronize() - assert watcher.branch == branch - grids[tag] = watcher.grid - _compare(out, op.reference(**args)) - assert grids["narrow"][1] == 1, grids - assert grids["wide"][1] > 1, grids - - -def _capture_args(seed, only_visible_blocks=False): - spec = _case( - 0, - 1, - layout="capture", - lengths=(8191,), - slots=8, - pages=CAPTURE_PAGES, - only_visible_blocks=only_visible_blocks, - context_rows=True, - ) - return _build(spec, seed) - - -def test_the_two_kernels_agree() -> 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; - asserted at the 8191-row capture shape. - """ - 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("context_rows", [False, True], ids=["row_wise", "query_tile"]) -def test_results_are_bitwise_reproducible(context_rows) -> None: - """Three identical calls, bitwise identical, in the padded mode so the - whole tensor is written; once per kernel.""" - args = _build(_case(300, 5, context_rows=context_rows), 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", "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: - args = _build(_case(64, 4, only_visible_blocks=True), seed=670) - 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_pdl_launch_is_bitwise_the_plain_one_and_reads_after_the_wait() -> None: - """`use_pdl` changes the launch, not the result; under graph replay a producer rewrites - q right before the op, and a read of it before the wait would see the previous - replay's.""" - args = _build(_case(4, 4, only_visible_blocks=True), seed=313) - visible = _visible(args) - plain = op(**args) - assert torch.equal(op(**{**args, "use_pdl": True})[visible], plain[visible]) - src = args["q"].clone() - - def step(): - args["q"].copy_(src * 1.0) - return op(**{**args, "use_pdl": True}) - - warmup = torch.cuda.Stream() - warmup.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(warmup): - step() - torch.cuda.current_stream().wait_stream(warmup) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - graph_out = step() - for trial in range(3): - src.copy_(_build(_case(4, 4, only_visible_blocks=True), seed=314 + trial)["q"]) - graph.replay() - torch.cuda.synchronize() - assert torch.equal(graph_out[visible], op(**args)[visible]), f"replay {trial}" - - def test_prefetch_reads_the_fresh_anchor_key_and_q_after_the_wait() -> None: """`prefetch` reads pages and keys before the PDL wait. Under graph replay a producer rewrites q and, at rows whose position ends a group, the last visible block's anchor @@ -409,177 +275,10 @@ def step(): assert torch.equal(graph_out[visible], op(**args)[visible]), f"replay {trial}" -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), - } - - -def test_wrong_variants_fail_the_gate() -> None: - """The default fp32 band is only a gate if a wrong answer fails it. - - At the target's 260-page width, each variant is rejected by `_compare` and - lands more than 1000x outside the band the correct result clears; measured - 3e5x to 4e6x. - """ - 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_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", - ), - "mismatched_dtypes": ( - lambda a: _mutated(a, index_cache=a["index_cache"].float()), - "dtypes must match", - ), - "compress_ratio_1": (lambda a: _mutated(a, compress_ratio=1), "compression ratio >= 2"), - "short_metadata": ( - lambda a: _mutated(a, query_positions=a["query_positions"][:4]), - "one entry per row", - ), -} - - -@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`. One case per refusal family: operand shape, block - table, dtypes, geometry scalars, metadata.""" - mutate, message = _REJECTIONS[case] - args = mutate(_build(_case(16, 2), seed=1701)) - with pytest.raises(ValueError, match=message): +def test_rejects_metadata_shorter_than_the_rows() -> None: + """The callable's own refusal, message as observed -- the reason this entry has no + `is_valid`: positions for fewer rows than `q` carries.""" + args = _build(_case(16, 2), seed=1701) + args["query_positions"] = args["query_positions"][:4] + with pytest.raises(ValueError, match="one entry per row"): op(**args) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_store.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_store.py index 85b8e69044c4..d6d7c5b03b47 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_store.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_store.py @@ -53,12 +53,9 @@ def _bits(tensor: torch.Tensor) -> torch.Tensor: def _hole(spec, table, allocated_cols): """One row that resolves to no slot: its (request, position).""" - kind = spec["hole"] - if kind == "unallocated_page": - assert int(table[0, allocated_cols[0]]) == -1 - return 0, allocated_cols[0] * spec["tokens_per_block"] - assert kind == "negative_request" - return -1, 0 + assert spec["hole"] == "unallocated_page", spec["hole"] + assert int(table[0, allocated_cols[0]]) == -1 + return 0, allocated_cols[0] * spec["tokens_per_block"] def _shared_page_planes(spec, extras: dict): @@ -181,61 +178,6 @@ def _small(seed: int, **overrides): 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 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_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 at the 32 MTP7 verify blocks (256 rows), replayed against new positions written into the same buffer.""" @@ -262,81 +204,13 @@ def test_cuda_graph_capture_and_replay() -> None: _compare((args["index_cache"], args["position_cache"]), ref) -def test_the_stored_key_is_raw_not_normalized() -> None: - """Discrimination against a Gemma-normed key, the decode kernel's form. - - A key stored normed double-normalizes every anchor `qsa_prefill_compress` - builds. Old receipt at its seeds: 3.0625 max abs apart; the floor states a - bound rather than one draw. - """ - args = _small(seed=66) - out = op(**args) - gen = torch.Generator(device=DEV).manual_seed(68) - weight = torch.randn(128, generator=gen, device=DEV) * 0.5 - x = args["token_k"].float() - normed = (x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) * (1 + weight)).to( - torch.bfloat16 - ) - page = int(args["block_table"][0, 0]) - stored = out[0][page, 0] - assert torch.equal(_bits(stored), _bits(args["token_k"][0])) - distance = (stored.float() - normed[0].float()).abs().max().item() - assert distance > 0.25, f"a normed key differs from the raw one by only {distance}" - - -# ------------------------------------------------------- loud rejections -# -# Each is a shape the kernel would otherwise read past or write wrongly; the -# entry refuses it before launching, so `is_valid` has nothing left to guard. -# One case per refusal family: metadata, row counts, the index plane, the -# position plane, the table. - - -def _rejected(args, match: str) -> None: - before = (args["index_cache"].clone(), args["position_cache"].clone()) - with pytest.raises(ValueError, match=match): - op(**args) - assert torch.equal(_bits(args["index_cache"]), _bits(before[0])) - assert torch.equal(args["position_cache"], before[1]) - - def test_rejects_metadata_of_unequal_length() -> None: + """A shape the kernel would otherwise read past; the entry refuses it before + launching, so `is_valid` has nothing left to guard.""" args = _small(seed=61) args["logical_positions"] = args["logical_positions"][:1] - _rejected(args, "equally long") - - -def test_rejects_a_single_token_k_row() -> None: - args = _small(seed=63) - args["token_k"] = args["token_k"][:1] - _rejected(args, "token_k") - - -def test_rejects_a_page_size_that_is_not_the_cache_s() -> None: - args = _small(seed=65) - args["tokens_per_block"] = 16 - _rejected(args, "index_cache") - - -def test_rejects_a_position_plane_narrower_than_the_coordinates() -> None: - args = _small(seed=67) - args["position_cache"] = torch.full((8, 32, 2), POS_FILL, dtype=torch.int32, device=DEV) - _rejected(args, "position_cache") - - -def test_rejects_a_one_dimensional_table() -> None: - args = _small(seed=69) - args["block_table"] = args["block_table"][0] - _rejected(args, "two-dimensional") - - -def test_huge_int64_position_is_skipped_not_truncated() -> None: - """`2^32 + 5` truncated to int32 would be position 5 of column 0; the - kernel keeps int64 and finds it past the table.""" - args = _small(seed=70, position_dtype=torch.int64) - args["logical_positions"][0] = 2**32 + 5 before = (args["index_cache"].clone(), args["position_cache"].clone()) - op(**args) - page = int(args["block_table"][0, 0]) - assert torch.equal(_bits(args["index_cache"][page, 5]), _bits(before[0][page, 5])) - assert torch.equal(args["position_cache"][page, 5], before[1][page, 5]) + with pytest.raises(ValueError, match="equally long"): + op(**args) + assert torch.equal(_bits(args["index_cache"]), _bits(before[0])) + assert torch.equal(args["position_cache"], before[1]) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store.py index 3a906789e826..caf83f14f1c7 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store.py @@ -50,25 +50,11 @@ def _bits(tensor: torch.Tensor) -> torch.Tensor: return tensor.contiguous().view(as_int) -def _hole(spec, table, live_requests, allocated_cols): - """One row that resolves to no slot, and the table edit it needs.""" - tpb, width, slots, pages = (spec[k] for k in ("tokens_per_block", "width", "slots", "pages")) - kind = spec["hole"] - if kind == "unallocated_page": - assert int(table[0, allocated_cols[0]]) == -1 - return 0, allocated_cols[0] * tpb - if kind == "page_past_pool": - table[live_requests] = -1 - table[live_requests, 0] = pages + 7 - return live_requests, 0 - if kind == "request_past_table": - return slots, 0 - if kind == "negative_request": - return -1, 0 - if kind == "negative_position": - return 0, -1 - assert kind == "column_past_table" - return 0, width * tpb +def _hole(spec, table, allocated_cols): + """One row that resolves to no slot: its (request, position).""" + assert spec["hole"] == "unallocated_page", spec["hole"] + assert int(table[0, allocated_cols[0]]) == -1 + return 0, allocated_cols[0] * spec["tokens_per_block"] def _build(spec, seed: int): @@ -91,7 +77,7 @@ def _build(spec, seed: int): allocated_cols.append(cols) assert used <= pages and len(segments) < slots, "pool or table too small for the batch" if spec["hole"] is not None: - request, position = _hole(spec, table, len(segments), allocated_cols) + request, position = _hole(spec, table, allocated_cols) requests.append(request) positions.append(position) @@ -105,26 +91,7 @@ def _build(spec, seed: int): ).to(device=DEV, dtype=spec["row_dtype"]) k = qkv[:, spec["heads_q"] : spec["heads_q"] + hkv] v = qkv[:, spec["heads_q"] + hkv : spec["heads_q"] + 2 * hkv] - if spec["values"] == "every_bf16": - patterns = torch.arange(-(2**15), 2**15, dtype=torch.int32).to(torch.int16) - patterns = patterns.view(torch.bfloat16).reshape(rows, hkv, head_dim).to(DEV) - k.copy_(patterns) - v.copy_(patterns.flip(0)) - else: - assert spec["values"] == "randn" - 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"] @@ -158,23 +125,13 @@ def test_certified_cells(cell, certifying): # noqa: F811 def _decode_args(seed: int, **overrides): - spec = op.CELLS[0].spec + """The decode cell's layout over a bf16 pool, without its skipped row.""" + spec = {**op.CELLS[0].spec, "hole": None} return _build({**spec, **overrides}, seed=seed) -def test_fp8_pool_saturates_an_out_of_range_row() -> None: - """1000.0 lands as e4m3's finite maximum, matching torch's own conversion.""" - args, _ = _decode_args(seed=4, cache_dtype=torch.float8_e4m3fn) - args["k"] = torch.full_like(args["k"], 1000.0) - out = op(**args) - page, slot = int(args["block_table"][0, 0]), 0 - stored = out[0][page, :, slot, :] - assert torch.equal(_bits(stored), _bits(args["k"][0].to(torch.float8_e4m3fn))) - assert float(stored[0, 0].float()) == 448.0 - - #: The cell of 32 MTP7 verify blocks (256 rows): the segments, pool and table it runs over. -_MTP7_32 = next(c.spec for c in op.CELLS if c.why.startswith("decode MTP7: 32 requests")) +_MTP7_32 = next(c.spec for c in op.CELLS if c.why.startswith("fp8 decode MTP7: 32 requests")) _MTP7_BATCH = {key: _MTP7_32[key] for key in ("segments", "pages", "slots")} @@ -204,37 +161,6 @@ def call(): _compare((args["k_cache"], args["v_cache"]), ref) -def test_op_rejects_operands_it_cannot_serve() -> None: - """The callable's own loud errors: why `is_valid` carries only the stride guard.""" - args, _ = _decode_args(seed=5) - row, cache = args["k"], args["k_cache"] - narrow = torch.zeros(row.shape[0], row.shape[1], 64, dtype=row.dtype, device=DEV) - cases = { - "QSA K/V inputs must use BF16 or FP16": dict(k=row.float(), v=row.float()), - "QSA K/V cache must use BF16, FP16, or FP8 E4M3": dict( - k_cache=cache.float(), v_cache=cache.float() - ), - "QSA K/V pairs must use matching dtypes": dict(v=args["v"].half()), - "QSA K/V cache token dimension does not match tokens_per_block": dict(tokens_per_block=16), - "QSA K/V inputs and caches have incompatible head geometry": dict(k=narrow, v=narrow), - "QSA K/V store request metadata must use integer storage": dict( - logical_positions=args["logical_positions"].float() - ), - "QSA K/V store metadata must have one entry per input row": dict( - request_indices=args["request_indices"][:1].contiguous() - ), - "QSA K/V store requires a two-dimensional int32 block table": dict( - block_table=args["block_table"].long() - ), - "QSA paged K/V store requires CUDA tensors": { - name: (t.cpu() if isinstance(t, torch.Tensor) else t) for name, t in args.items() - }, - } - for message, override in cases.items(): - with pytest.raises(ValueError, match=message): - op(**{**args, **override}) - - def test_guard_refuses_strided_row_metadata() -> None: """`is_valid`'s one guard, after the hazard it exists for. @@ -243,44 +169,22 @@ def test_guard_refuses_strided_row_metadata() -> None: buffer is read through its base pointer: accepted, and rows land elsewhere. """ args, pool = _decode_args(seed=10) - requests, positions = args["request_indices"], args["logical_positions"] + requests = args["request_indices"] rows = requests.numel() wide_requests = torch.full((rows, 2), rows - 1, dtype=requests.dtype, device=DEV) wide_requests[:, 0] = requests - wide_positions = torch.full((rows, 2), 31, dtype=positions.dtype, device=DEV) - wide_positions[:, 0] = positions ref = op.reference(**args) + strided = {**args, "request_indices": wide_requests[:, 0]} - for override in ( - dict(request_indices=wide_requests[:, 0]), - dict(logical_positions=wide_positions[:, 0]), - ): - pool.fill_(FILL) - out = op(**{**args, **override}) # unguarded: accepted - assert not all(torch.equal(_bits(got), _bits(want)) for got, want in zip(out, ref)), ( - "strided metadata was expected to place rows elsewhere" - ) - - pool.fill_(FILL) - before = pool.clone() - with pytest.raises(AssertionError, match="must be a contiguous 1-D tensor"): - with validating(op): - op(**{**args, **override}) - assert torch.equal(_bits(pool), _bits(before)), "a refused call still wrote" - - -def test_int64_position_past_int32_is_truncated_not_skipped() -> None: - """Hard-limit probe: the kernel casts each position to int32. - - `2^32 + 5` should resolve to a page column far past the table and be - skipped; truncated, it is position 5, a real slot, and it is written. - """ - args, pool = _decode_args(seed=11) - args["logical_positions"][0] = 2**32 + 5 - ref = op.reference(**args) - out = op(**args) - page = int(args["block_table"][0, 0]) - assert torch.equal(_bits(out[0][page, :, 5, :]), _bits(args["k"][0].to(out[0].dtype))), ( - "position 2^32 + 5 was expected to land on slot 5 of page column 0" + pool.fill_(FILL) + out = op(**strided) # unguarded: accepted + assert not all(torch.equal(_bits(got), _bits(want)) for got, want in zip(out, ref)), ( + "strided metadata was expected to place rows elsewhere" ) - assert not torch.equal(_bits(out[0]), _bits(ref[0])), "the reference skips that row" + + pool.fill_(FILL) + before = pool.clone() + with pytest.raises(AssertionError, match="must be a contiguous 1-D tensor"): + with validating(op): + op(**strided) + assert torch.equal(_bits(pool), _bits(before)), "a refused call still wrote" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store_nvfp4.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store_nvfp4.py index 8d939576ef43..67f8d87f6179 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store_nvfp4.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store_nvfp4.py @@ -49,42 +49,17 @@ def _compare(out, ref) -> None: ) -def _hole(spec, table, live_requests, allocated_cols): - """One row that resolves to no slot, and the table edit it needs.""" - tpb, width, slots, pages = (spec[k] for k in ("tokens_per_block", "width", "slots", "pages")) - kind = spec["hole"] - if kind == "unallocated_page": - assert int(table[0, allocated_cols[0]]) == -1 - return 0, allocated_cols[0] * tpb - if kind == "page_past_pool": - table[live_requests] = -1 - table[live_requests, 0] = pages + 7 - return live_requests, 0 - if kind == "request_past_table": - return slots, 0 - if kind == "negative_request": - return -1, 0 - if kind == "negative_position": - return 0, -1 - assert kind == "column_past_table" - return 0, width * tpb +def _hole(spec, table, allocated_cols): + """One row that resolves to no slot: its (request, position).""" + assert spec["hole"] == "unallocated_page", spec["hole"] + assert int(table[0, allocated_cols[0]]) == -1 + return 0, allocated_cols[0] * spec["tokens_per_block"] def _fill_rows(spec, k, v, gen): - """Row values: random with per-block magnitudes over 2^-10..2^5, every finite bf16, - or edge blocks (zero, -0.0, an outlier).""" + """Row values: random with per-block magnitudes over 2^-10..2^5, plus edge blocks + (zero, -0.0, an outlier) when the cell asks for them.""" hkv, head_dim = spec["hkv"], spec["head_dim"] - if spec["values"] == "finite_bf16": - bits = torch.arange(-(2**15), 2**15, dtype=torch.int32).to(torch.int16) - finite = bits.view(torch.bfloat16) - finite = finite[torch.isfinite(finite.float())] # 65280 values - padded = torch.zeros(k.numel(), dtype=torch.bfloat16) - assert finite.numel() <= padded.numel() - padded[: finite.numel()] = finite - k.copy_(padded.reshape(k.shape).to(DEV)) - order = torch.randperm(padded.numel(), generator=gen) - v.copy_(padded[order].reshape(v.shape).to(DEV)) - return rows = k.shape[0] for t in (k, v): base = torch.randn(rows, hkv, head_dim // 16, 16, generator=gen) @@ -155,7 +130,7 @@ def _build(spec, seed: int, pool=None): allocated_cols.append(cols) assert used <= pages and len(segments) <= slots, "pool or table too small for the batch" if spec["hole"] is not None: - request, position = _hole(spec, table, len(segments), allocated_cols) + request, position = _hole(spec, table, allocated_cols) requests.append(request) positions.append(position) @@ -198,7 +173,7 @@ def _pools_equal(a, b) -> bool: return all(torch.equal(x, y) for x, y in zip(a, b)) -def _graph_replay(spec, seed: int) -> None: +def _graph_replay(spec, seed: int, graph_batches) -> None: """Each G: capture one step of the first G requests, replay the next step. The next step lives in the captured buffers: new k/v rows, new pages in the @@ -210,7 +185,7 @@ def _graph_replay(spec, seed: int) -> None: steps = {n for _, n in spec["segments"]} assert len(steps) == 1, "a graph cell's requests all carry the same row count" step = steps.pop() - for g in spec["graph_batches"]: + for g in graph_batches: sub = {**spec, "segments": spec["segments"][:g]} args, _, _ = _build(sub, seed=seed + g, pool=pool) following, _, _ = _build(sub, seed=seed + 1000 + g, pool=pool) @@ -356,9 +331,6 @@ def _thop_equality(spec, seed: int) -> None: def test_certified_cells(cell, certifying): # noqa: F811 certifying(op) seed = zlib.crc32(cell.why.encode()) - if cell.spec["graph_batches"]: - _graph_replay(cell.spec, seed=seed) - return if cell.spec["thop_append"]: _thop_equality(cell.spec, seed=seed) return @@ -375,91 +347,23 @@ def test_certified_cells(cell, certifying): # noqa: F811 _assert_other_roles_untouched(cell.spec, pool) -def _specialization_class(value: int) -> str: - """Triton's class for a plain-int kernel argument: 1, a multiple of 16, or other.""" - if value == 1: - return "one" - return "multiple_of_16" if value % 16 == 0 else "other" - - -def test_cells_reach_every_served_specialization_class() -> None: - """The pool's slot count and the table's height each select a compiled variant - by class. Slot count other or a multiple of 16 against table height 1, other - or a multiple of 16 all run, eagerly.""" - reached = { - (_specialization_class(c.spec["pages"]), _specialization_class(c.spec["slots"])) - for c in op.CELLS - if not c.spec["graph_batches"] - } - wanted = { - (pages, height) - for pages in ("other", "multiple_of_16") - for height in ("one", "other", "multiple_of_16") - } - assert wanted <= reached, f"unreached: {sorted(wanted - reached)}" +def test_cuda_graph_replay_matches_eager() -> None: + """Decode MTP under a graph: 8 requests x 8 rows in the prefill cell's trunk layout, + captured once, replayed on the next step's rows, pages and positions.""" + prefill = next(c.spec for c in op.CELLS if c.why.startswith("prefill:")) + spec = {**prefill, "segments": tuple((29 + 61 * i, 8) for i in range(8)), "pages": 256} + _graph_replay(spec, seed=8, graph_batches=(8,)) def _decode_args(seed: int, **overrides): return _build({**op.CELLS[0].spec, **overrides}, seed=seed) -def test_rounding_rivals_are_rejected() -> None: - """The bitwise gate separates the engine's rounding from near misses on the - every-finite-bf16 rows: e2m1 ties rounded away from zero instead of to even, - and V's scales laid out token-major like K's, each move bytes the gate sees.""" - from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention import ( # noqa: E501 - qsa_paged_kv_store_nvfp4 as entry, - ) - - spec = next(c.spec for c in op.CELLS if c.spec["values"] == "finite_bf16") - args, _, _ = _build(spec, seed=21) - out = op(**args) - _compare(out, op.reference(**args)) - - ties_up = tuple((threshold, True) for threshold, _ in entry._E2M1_STEPS) - saved = entry._E2M1_STEPS - try: - entry._E2M1_STEPS = ties_up - rival = op.reference(**args) - finally: - entry._E2M1_STEPS = saved - moved = int((rival[0] != out[0]).sum()) - assert moved > 0, "ties away from zero must move some K bytes on every-finite-bf16 rows" - - # V's scale layout, on random rows (the every-finite rows saturate V's scales). - spec = next(c.spec for c in op.CELLS if c.why.startswith("prefill:")) - args, _, _ = _build(spec, seed=22) - out = op(**args) - page = int(args["block_table"][0, 0]) - groups = spec["head_dim"] // 16 - interleaved = out[3][page].reshape(spec["hkv"], -1) - token_major = ( - interleaved.reshape(spec["hkv"], spec["tokens_per_block"] // 4, groups, 4) - .permute(0, 1, 3, 2) - .reshape(spec["hkv"], -1) - ) - assert not torch.equal(interleaved, token_major), "the V scale layouts must be told apart" - - -def test_op_rejects_operands_it_cannot_serve() -> None: - """The callable's own loud errors: why the entry carries no `is_valid`.""" +def test_op_rejects_strided_request_indices() -> None: + """The callable's own loud error -- why the entry carries no `is_valid` -- for the + metadata a flat-pointer kernel would otherwise read at the wrong rows.""" args, _, _ = _decode_args(seed=5) - cache, scales = args["k_cache"], args["k_scale_cache"] - cases = { - "k and v must both be BF16 or both FP16": dict(k=args["k"].float(), v=args["v"].float()), - "k_cache must be uint8": dict(k_cache=cache.view(torch.int8)), - "k_scale_cache must be uint8": dict(k_scale_cache=scales[..., :8]), - "tokens_per_block must be a positive multiple of 4": dict(tokens_per_block=30), - "block_table must be a 2-D int32 tensor": dict(block_table=args["block_table"].long()), - "request_indices must be contiguous": dict( - request_indices=torch.stack([args["request_indices"]] * 2, 1)[:, 0] - ), - "kv_scale_orig_quant must be a contiguous fp32": dict( - kv_scale_orig_quant=args["kv_scale_orig_quant"][:1] - ), - "every operand must be on k's CUDA device": dict(block_table=args["block_table"].cpu()), - } - for message, override in cases.items(): - with pytest.raises(ValueError, match=message): - with validating(op): - op(**{**args, **override}) + strided = torch.stack([args["request_indices"]] * 2, 1)[:, 0] + with pytest.raises(ValueError, match="request_indices must be contiguous"): + with validating(op): + op(**{**args, "request_indices": strided}) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa.py index b44ab112e3dc..28ed2deb6c8a 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa.py @@ -20,9 +20,6 @@ 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, -) 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, ) @@ -47,7 +44,7 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: flip, since bf16's step is 2 of these units at the bottom of a binade -- and at most 0.62 ulp RMS: 2x and 3.2x margin. Six wrong references (no scale, interleaved heads, padding attended, identity pages, K/V - swapped, page/slot transposed) are each more than 0.1 off and rejected. + swapped, page/slot transposed) were each measured more than 0.1 off. """ assert out.dtype == ref.dtype, (out.dtype, ref.dtype) assert out.shape == ref.shape, (out.shape, ref.shape) @@ -57,10 +54,6 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: DEV = "cuda" -ENGINE_MAX_SEQ_LEN = 8192 -FP8 = torch.float8_e4m3fn -#: The e4m3fn NaN bit pattern (and, with the sign bit, 0xFF). -E4M3_NAN = 0x7F _KERNELS = { "_qsa_paged_sparse_gqa_kernel": "prefill", @@ -95,26 +88,6 @@ def launches(monkeypatch): # -------------------------------------------------------------------------- -def _segments(spec): - """`(first position, rows)` per request in row order, and a generation mask.""" - if "mixed" not in spec: - segments = list(spec["segments"]) - return segments, torch.zeros(sum(n for _, n in segments), dtype=torch.bool, device=DEV) - requests, context_lens = spec["mixed"] - # Chunk 0 starts at position 0, so one launch spans rows that see one token - # and rows at the full budget; later chunks are chunked-prefill continuations. - segments = [ - (0 if i == 0 else (907 * i) % (ENGINE_MAX_SEQ_LEN - n + 1), n) - for i, n in enumerate(context_lens) - ] - # MTP: two rows per generation request, positions across the budget both ways. - segments += [(100 + (499 * g) % 8000, 2) for g in range(requests - len(context_lens))] - generation = torch.cat( - [torch.full((n,), i >= len(context_lens), device=DEV) for i, (_, n) in enumerate(segments)] - ) - return segments, generation - - def _selected_tokens(positions, seed, compress_ratio, token_topk): """QSA's expansion layout, built in one pass. @@ -156,17 +129,14 @@ def _build(spec, seed): spec["head_dim"], spec["q_heads"], ) - segments, generation = _segments(spec) + segments = list(spec["segments"]) positions = torch.cat([torch.arange(s, s + n, device=DEV) for s, n in segments]) req = torch.cat([torch.full((n,), i, device=DEV) for i, (_, n) in enumerate(segments)]) rows, slots = positions.numel(), spec["slots"] assert len(segments) <= slots slot_of = torch.randperm(slots, generator=torch.Generator().manual_seed(seed))[: len(segments)] - # `min_seq_len` maps pages past the last query row's position, so a probe - # can plant a token that is resident but not causally visible. - floor = spec.get("min_seq_len", 0) - mapped = [math.ceil(max(s + n, floor) / tpb) for s, n in segments] + mapped = [math.ceil((s + n) / tpb) for s, n in segments] assert max(mapped) <= spec["pages"] pool_pages = sum(mapped) + 7 page_ids = torch.randperm(pool_pages, generator=gen, device=DEV).to(torch.int32) @@ -178,35 +148,21 @@ def _build(spec, seed): pool = torch.randn(pool_pages, 2, kv, tpb, d, generator=gen, device=DEV).to(spec["cache_dtype"]) k_cache, v_cache = pool[:, 0], pool[:, 1] - if spec["cache_layout"] == "contiguous": - k_cache, v_cache = k_cache.contiguous(), v_cache.contiguous() q_dtype = spec["q_dtype"] - if spec["q_layout"] in ("qkv_slice", "fused_row"): - row_heads = q_heads + 2 * kv if spec["q_layout"] == "qkv_slice" else 55 - qkv = torch.randn(rows, row_heads, d, generator=gen, device=DEV).to(q_dtype) - q = qkv[:, :q_heads] - else: - q = torch.randn(rows, q_heads, d, generator=gen, device=DEV).to(q_dtype) + row_heads = q_heads + 2 * kv if spec["q_layout"] == "qkv_slice" else 55 + qkv = torch.randn(rows, row_heads, d, generator=gen, device=DEV).to(q_dtype) + q = qkv[:, :q_heads] selected = _selected_tokens(positions, seed + 3, spec["compress_ratio"], spec["token_topk"]) - if "width" in spec: - selected = selected[:, : spec["width"]].contiguous() request_indices = slot_of.to(DEV)[req].to(spec["request_dtype"]) - if spec.get("holes"): - selected[1] = -1 # an all-padding row - request_indices[2] = slots + 67 # a request the table does not have - block_table[:, 0] = -1 # every request's first page unmapped - gate = None if spec.get("gate") == "fused_row": # The gate columns of the target's fused projection row (55 heads of # D): heads 28..51, after q, k and v. fused = torch.randn(rows, 55 * d, generator=gen, device=DEV).to(q_dtype) gate = fused[:, 28 * d : (28 + q_heads) * d] - elif spec.get("gate") == "heads": - gate = torch.randn(rows, q_heads, d, generator=gen, device=DEV).to(q_dtype) args = dict( q=q, @@ -222,7 +178,7 @@ def _build(spec, seed): args.update(query_positions=positions, compress_ratio=spec["compress_ratio"]) if gate is not None: args.update(output_gate=gate) - return args, generation + return args def _bound(positions, spec): @@ -242,9 +198,9 @@ def _is_prefix_packed(table): def test_certified_cells(cell, certifying, launches) -> None: # noqa: F811 certifying(op) spec = cell.spec - args, _ = _build(spec, seed=sum(map(ord, cell.why))) + args = _build(spec, seed=sum(map(ord, cell.why))) rows = args["q"].shape[0] - if spec["prefill_pair"] and not spec.get("holes"): + if spec["prefill_pair"]: # The precondition the op cannot check: front-packed, and exactly as # many valid entries as the prefill kernel will scan. selected = args["selected_tokens"] @@ -261,11 +217,7 @@ def test_certified_cells(cell, certifying, launches) -> None: # noqa: F811 _compare(out, op.reference(**args)) -# -------------------------------------------------------------------------- -# the prefill kernel's bound -# -------------------------------------------------------------------------- - -#: The first cell is this checkpoint's own decode geometry; the probes below +#: The first cell is this checkpoint's own decode geometry; the tests below #: start from it. CKPT = op.CELLS[0].spec assert (CKPT["q_heads"], CKPT["kv_heads"], CKPT["head_dim"]) == (24, 2, 256) @@ -275,328 +227,24 @@ def _with(**overrides): return {**CKPT, **overrides} -def _first_unscanned_column(positions, tile=64): - """First column of each row the prefill kernel never reads (bound rounded - up to its 64-wide tile); none when that lands at or past the row's width.""" - final_topk = CKPT["token_topk"] + CKPT["compress_ratio"] - 1 - bound = _bound(positions, CKPT).clamp(max=final_topk) - return ((bound + tile - 1) // tile) * tile, final_topk - - -def _plant_past_the_bound(args): - """Put each row's own newest token in its first unscanned column.""" - positions = args["query_positions"] - first, final_topk = _first_unscanned_column(positions) - gap_rows = (first < final_topk).nonzero().flatten() - assert gap_rows.numel() > 0, "no row has an unscanned column to plant into" - planted = args["selected_tokens"].clone() - planted[gap_rows, first[gap_rows]] = positions[gap_rows].to(torch.int32) - return planted, gap_rows - - -def _sixty_four_token_rows(rows, seed): - """Rows seeing 64 tokens: the bound is exactly one tile and column 64 the - first unread one. It is planted with token 200 -- resident (each request - maps 256 tokens) and not otherwise selected.""" - spec = _with(segments=((63, 1),) * rows, slots=rows + 8, min_seq_len=256, prefill_pair=True) - args, _ = _build(spec, seed) - packed = args["selected_tokens"] - planted = packed.clone() - planted[:, 64] = 200 - return {**args, "selected_tokens": planted}, packed - - -def test_prefill_kernel_stops_at_the_causal_prefix(launches) -> None: - """An index past the bound is silently not attended to. - - The op matches the reference computed WITHOUT the planted index, and the - planted index moves the correct answer by more than 0.1, so the first - comparison is not vacuous. - """ - args, packed = _sixty_four_token_rows(300, seed=7000) - out = op(**args) - assert [name for name, _ in launches] == ["prefill"] - without_it = op.reference(**{**args, "selected_tokens": packed}) - _compare(out, without_it) - with_it = op.reference(**args) - assert (with_it.float() - without_it.float()).abs().max().item() > 0.1 - - -def test_split_k_reads_the_whole_row(launches) -> None: - """Below the program threshold the same operands attend to every column.""" - args, packed = _sixty_four_token_rows(100, seed=7100) - out = op(**args) - assert [name for name, _ in launches][0] == "splitk" - _compare(out, op.reference(**args)) - without_it = op.reference(**{**args, "selected_tokens": packed}) - assert (out.float() - without_it.float()).abs().max().item() > 0.1 - - -def test_prefill_bound_at_8191_rows() -> None: - """The bound is the same at 8191 rows, and reaches much further: a row - seeing 64 tokens is scanned for 64 of 2051 columns.""" - args, _ = _build(_with(segments=((0, 8191),), slots=8, prefill_pair=True), seed=11000) - planted, _ = _plant_past_the_bound(args) - out = op(**{**args, "selected_tokens": planted}) - without_it = op.reference(**args) - _compare(out, without_it) - with_it = op.reference(**{**args, "selected_tokens": planted}) - assert (with_it.float() - without_it.float()).abs().max().item() > 0.1 - - -def test_prefill_bound_reaches_the_generation_rows() -> None: - """A mixed launch's generation rows ride the prefill kernel, bound and all.""" - args, generation = _build( - _with(mixed=(32, (900, 200, 150, 28)), slots=32, prefill_pair=True), seed=22000 - ) - planted, gap_rows = _plant_past_the_bound(args) - assert bool(generation[gap_rows].any()), "no generation row has an unscanned column" - out = op(**{**args, "selected_tokens": planted}) - without_it = op.reference(**args) - _compare(out, without_it) - with_it = op.reference(**{**args, "selected_tokens": planted}) - generation_gaps = gap_rows[generation[gap_rows]] - moved = (with_it[generation_gaps].float() - without_it[generation_gaps].float()).abs().max() - assert moved.item() > 0.1, f"the planted indices must matter to the generation rows: {moved}" - - -# -------------------------------------------------------------------------- -# semantics a cell cannot state -# -------------------------------------------------------------------------- - - -def test_repeated_index_is_attended_twice() -> None: - """Selected indices are not deduplicated; each column is its own key.""" - args, _ = _build(_with(segments=((900, 1),), prefill_pair=False), seed=8700) - selected = torch.full((1, 64), -1, dtype=torch.int32, device=DEV) - selected[0, :3] = torch.tensor([5, 5, 40], dtype=torch.int32, device=DEV) - twice = op(**{**args, "selected_tokens": selected}) - _compare(twice, op.reference(**{**args, "selected_tokens": selected})) - once = selected.clone() - once[0, 1] = -1 - single = op.reference(**{**args, "selected_tokens": once}) - assert (twice.float() - single.float()).abs().max().item() > 1e-2 - - def test_calls_are_bitwise_deterministic() -> None: """Including a 32-split launch whose merge is a second kernel.""" - args, _ = _build(_with(segments=tuple((4000 + 37 * i, 2) for i in range(4))), seed=8800) + spec = _with(segments=tuple((4000 + 37 * i, 2) for i in range(4)), gate=None) + args = _build(spec, seed=8800) assert torch.equal(op(**args), op(**args)) -@pytest.mark.parametrize("rows", [1, 128]) -def test_output_gate_in_the_merge_is_the_standalone_gate(rows) -> None: - """The gate carried by the merge equals the ungated call followed by the - standalone gate kernel, bit for bit, and the gate moves the output.""" - spec = _with(segments=((4000, 1),) * rows, slots=rows + 32, q_layout="fused_row") - spec["gate"] = "fused_row" - args, _ = _build(spec, seed=8900 + rows) - gate = args.pop("output_gate") - assert op.fuses_output_gate(args["q"], args["k_cache"], args["selected_tokens"], gate) - ungated = op(**args) - gated = op(**args, output_gate=gate) - standalone = fused_sigmoid_mul_inplace(ungated.reshape(rows, -1).clone(), gate) - assert torch.equal(gated.reshape(rows, -1), standalone) - assert (gated.float() - ungated.float()).abs().max().item() > 0.1 - - -def test_output_gate_needs_a_merge() -> None: - """One split (300 rows) and the prefill kernel have no merge to carry the - gate: `fuses_output_gate` says so and the call refuses it.""" - for spec in ( - _with(segments=((600, 1),) * 300, slots=384), - _with(segments=((2032, 272),), prefill_pair=True), - ): - spec["gate"] = "heads" - args, _ = _build(spec, seed=9200) - gate = args["output_gate"] - assert not op.fuses_output_gate(args["q"], args["k_cache"], args["selected_tokens"], gate) - with pytest.raises(ValueError, match="output-gate fusion requires a compatible split-K"): - 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"), - "int64_block_table": ( - lambda a: {**a, "block_table": a["block_table"].long()}, - "QSA sparse GQA requires a two-dimensional int32 block table", - ), - "negative_scale": ( - lambda a: {**a, "softmax_scale": -1.0}, - "QSA sparse GQA softmax scale must be finite and positive", - ), - "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. - One per guard family: operand dtype, metadata dtype, scalar value, head - geometry, the prefill pair and the device. 192 is a legal head size - elsewhere in this engine and not here.""" - args, _ = _build(_with(segments=((600, 1),) * 4), seed=9100) +def test_rejects_positions_without_a_compress_ratio() -> None: + """The op's own loud refusal, verbatim -- the reason `is_valid` is empty: half of + the prefill pair.""" + args = _build(_with(segments=((600, 1),) * 4, gate=None), seed=9100) args["selected_tokens"] = args["selected_tokens"][:, :64].contiguous() - mutate, message = _REJECTIONS[case] + args["query_positions"] = torch.zeros(4, dtype=torch.int64, device=DEV) + message = "QSA causal-prefix bounds require positions and compression ratio" with pytest.raises(ValueError, match=re.escape(message)): - op(**mutate(args)) - - -def test_selection_width_is_checked_on_the_prefill_kernel_only() -> None: - """A width the compress ratio cannot produce: refused by the prefill - kernel, accepted (the pair is unused) by split-K.""" - args, _ = _build(_with(segments=((200, 272),), prefill_pair=True), seed=9300) - args["selected_tokens"] = args["selected_tokens"][:, :64].contiguous() - with pytest.raises(ValueError, match="incompatible with its compression ratio"): op(**args) - - args, _ = _build(_with(segments=((200, 100),), prefill_pair=True), seed=9400) - args["selected_tokens"] = args["selected_tokens"][:, :64].contiguous() - _compare(op(**args), op.reference(**args)) - - -# -------------------------------------------------------------------------- -# an fp8 pool is read at face value -# -------------------------------------------------------------------------- - - -def _every_finite_e4m3_code(): - """The 256 e4m3fn codes as bytes, the two NaN codes replaced by +0.""" - codes = torch.arange(256, dtype=torch.int32, device=DEV).to(torch.uint8) - return torch.where((codes & 0x7F) == E4M3_NAN, torch.zeros_like(codes), codes) - - -def _two_token_rows(kernel, seed): - """Rows at position 1 of a request each, on an fp8 pool: every row's - selection is tokens 0 and 1, both on the request's first page.""" - rows = 300 if kernel == "prefill" else 12 - spec = _with( - segments=((1, 1),) * rows, - slots=rows + 8, - prefill_pair=kernel == "prefill", - cache_dtype=FP8, - ) - args, _ = _build(spec, seed) - pages = args["block_table"][args["request_indices"].long(), 0].long() - return args, pages - - -@pytest.mark.parametrize("kernel", ["prefill", "splitk"]) -def test_fp8_values_are_read_at_face_value(kernel, launches) -> None: - """Token 0's V row holds every finite code, token 1's is zero, and q is - zero, so each output row is exactly half of token 0's row: no scale, and - the widening to bf16 is exact for every code, subnormals included.""" - args, pages = _two_token_rows(kernel, seed=12000) - codes = _every_finite_e4m3_code() - values = args["v_cache"].view(torch.uint8) - values[pages, :, 0, :] = codes - values[pages, :, 1, :] = 0 - args["q"].zero_() - out = op(**args) - assert launches[0][0] == kernel - expected = codes.view(FP8).to(torch.bfloat16) / 2 - assert torch.equal(out, expected.expand_as(out)) - _compare(out, op.reference(**args)) - - -@pytest.mark.parametrize("kernel", ["prefill", "splitk"]) -def test_fp8_keys_are_read_at_face_value(kernel, launches) -> None: - """Token 0's K row holds every finite code, token 1's is zero, and each - (row, head) queries one element of token 0's row with a one-hot power of - two sized so the logit lands near +-8. Every product is exact, and each - output row is token 0's softmax weight for one code. A reference that - flushed K's 14 subnormal codes to zero moves their weights to 0.5 and is - rejected, so the gate sees a misread of a single code.""" - args, pages = _two_token_rows(kernel, seed=12100) - codes = _every_finite_e4m3_code() - keys, values = args["k_cache"].view(torch.uint8), args["v_cache"].view(torch.uint8) - keys[pages, :, 0, :] = codes - keys[pages, :, 1, :] = 0 - values[pages, :, 0, :] = torch.tensor(1.0, device=DEV).to(FP8).view(torch.uint8) - values[pages, :, 1, :] = 0 - q = args["q"] - rows, heads, dim = q.shape - assert rows * heads >= dim, "every code must be read by some (row, head)" - magnitude = codes.view(FP8).float().abs() - # |q . k| * softmax_scale = 128 / 16 = 8 up to the power-of-two rounding. - scale = torch.where(magnitude > 0, torch.exp2(torch.round(torch.log2(128.0 / magnitude))), 1.0) - column = (torch.arange(rows * heads, device=DEV) % dim).reshape(rows, heads, 1) - q.zero_() - q.scatter_(2, column, scale[column].to(q.dtype)) - out = op(**args) - assert launches[0][0] == kernel - _compare(out, op.reference(**args)) - - subnormal = ((codes & 0x78) == 0) & ((codes & 0x07) != 0) - assert int(subnormal.sum()) == 14 - flushed = args["k_cache"].clone() - flushed.view(torch.uint8)[pages, :, 0, :] = torch.where(subnormal, codes & 0x80, codes) - wrong = op.reference(**{**args, "k_cache": flushed}) - with pytest.raises(AssertionError): - _compare(out, wrong) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_decode.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_decode.py index 99060ead085b..7d24c80e8471 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_decode.py @@ -2,18 +2,14 @@ # SPDX-License-Identifier: Apache-2.0 """GPU test for the qsa_paged_sparse_gqa_decode entry: cells against the composition entry's fp32 reference in that entry's band (4 ulp of each head row's magnitude, 2 ulp -relative RMS), agreement with the composition entry (`qsa_paged_sparse_gqa`, which this -replaces at decode) in the same band, stale-read checks under graph replay with PDL, and -the guards.""" +relative RMS), the selection's read-after-wait under graph replay with PDL, and the row +guard.""" import zlib import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_sparse_gqa import ( # noqa: E501 - qsa_paged_sparse_gqa, -) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_sparse_gqa_decode import ( # noqa: E501 qsa_paged_sparse_gqa_decode as op, ) @@ -71,21 +67,6 @@ def _build(spec, seed): return [q, k_cache, v_cache, block_table, selected, request_indices, TPB, D**-0.5, gate] -def _kw(args): - names = ( - "q", - "k_cache", - "v_cache", - "block_table", - "selected_tokens", - "request_indices", - "tokens_per_block", - "softmax_scale", - "output_gate", - ) - return dict(zip(names, args)) - - @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name certifying(op) @@ -96,15 +77,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0) -def test_matches_the_composition() -> None: - """One representative gated decode row; the composition is certified on its own.""" - args = _build(dict(pos=(8700,), gated=True), seed=91) - got = op(*args) - want = qsa_paged_sparse_gqa(**_kw(args)) - assert_within_ulp(got, want, element_ulp=4.0, rms_ulp=2.0) - print(f"\n bitwise the composition's {(got == want).float().mean().item():.4f}") - - def _replays_like_eager(step, refresh) -> None: warmup = torch.cuda.Stream() warmup.wait_stream(torch.cuda.current_stream()) @@ -121,21 +93,6 @@ def _replays_like_eager(step, refresh) -> None: assert torch.equal(graph_out, step()), f"replay {trial}" -def test_reads_q_and_the_gate_fresh_under_graph_replay() -> None: - """A plain launch rewrites q and the gate (one fused row) inside the graph right before - the op. The op reads both before its waits, which is sound only because such a producer - releases the op at its end: this pins that contract.""" - args = _build(dict(pos=(8700,), gated=True), seed=17) - fused = args[0]._base if args[0]._base is not None else args[0] - src = fused.clone() - - def step(): - fused.copy_(src * 1.0) - return op(*args) - - _replays_like_eager(step, lambda: src.copy_(torch.randn(src.shape, device=DEV).to(src.dtype))) - - def test_reads_the_selection_after_the_wait_under_graph_replay() -> None: """The target's top-k runs inside the graph right before the op: a PDL producer that releases the op right after its own wait, before it writes the selection. A read of the @@ -154,10 +111,8 @@ def step(): _replays_like_eager(step, lambda: scores.copy_(torch.rand(scores.shape, device=DEV))) -def test_guards() -> None: +def test_guard_refuses_rows_past_max_rows() -> None: args = _build(dict(pos=(8700,), gated=True), seed=5) many = [a.expand(300, *a.shape[1:]) if i in (0, 4, 5) else a for i, a in enumerate(args[:6])] with validating(op), pytest.raises(AssertionError): op(*many, *args[6:8], None) # 300 rows: past MAX_ROWS - with validating(op), pytest.raises(AssertionError): - op(*args[:8], args[8].float()) # gate not in q's dtype diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4.py index fcb9ab66568e..ab4313b61e21 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4.py @@ -80,23 +80,6 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0) -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"] - segments = [ - (0 if i == 0 else (907 * i) % (ENGINE_MAX_SEQ_LEN - n + 1), n) - for i, n in enumerate(context_lens) - ] - 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: chosen complete groups, the trailing partial group, -1.""" block_topk = token_topk // compress_ratio @@ -133,7 +116,7 @@ def _build(spec, seed): """Real paged state for one spec, keyed by the op's parameter names.""" gen = torch.Generator(device=DEV).manual_seed(seed) tpb, kv, d, q_heads = (spec[k] for k in ("tokens_per_block", "kv_heads", "head_dim", "q_heads")) - segments, generation = _segments(spec) + segments = list(spec["segments"]) positions = torch.cat([torch.arange(s, s + n, device=DEV) for s, n in segments]) req = torch.cat([torch.full((n,), i, device=DEV) for i, (_, n) in enumerate(segments)]) rows, slots = positions.numel(), spec["slots"] @@ -161,20 +144,11 @@ def _build(spec, seed): block_table[slot, :count] = page_ids[cursor : cursor + count] cursor += count - 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) + assert spec["q_layout"] == "qkv_slice", spec["q_layout"] + qkv = torch.randn(rows, q_heads + 2 * kv, d, generator=gen, device=DEV).to(spec["q_dtype"]) + q = qkv[:, :q_heads] selected = _selected_tokens(positions, seed + 3, spec["compress_ratio"], spec["token_topk"]) - 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 - request_indices[2] = slots + 67 - block_table[:, 0] = -1 args = dict( q=q, @@ -193,7 +167,7 @@ def _build(spec, seed): ) if spec["prefill_pair"]: args.update(query_positions=positions, compress_ratio=spec["compress_ratio"]) - return args, generation + return args def _bf16_twin(args): @@ -274,8 +248,8 @@ def _capture_once(shape, seed, launches, where) -> None: """One call of `shape` captured, a second state's operands written in place, replayed: the replay must be bitwise the eager call on the second state.""" where = f"{where}, [{shape['slots']}, {shape['pages']}] table" - args, _ = _build(shape, seed) - fresh, _ = _build(shape, seed + 50) + args = _build(shape, seed) + fresh = _build(shape, seed + 50) launches.clear() op(**args) assert [name for name, _ in launches] == list(shape["launches"]), where @@ -303,9 +277,9 @@ def test_certified_cells(cell, certifying, launches) -> None: # noqa: F811 if "capture" in spec: _capture_sweep(spec, seed=sum(map(ord, cell.why)), launches=launches) return - args, _ = _build(spec, seed=sum(map(ord, cell.why))) + args = _build(spec, seed=sum(map(ord, cell.why))) rows = args["q"].shape[0] - if spec["prefill_pair"] and not spec.get("holes"): + if spec["prefill_pair"]: selected = args["selected_tokens"] assert bool(_is_prefix_packed(selected).all()) valid = (selected >= 0).sum(1) @@ -330,81 +304,18 @@ def _with(**overrides): return {**CKPT, **overrides} -def test_gate_discriminates_wrong_variants() -> None: - """References wrong one way each -- V's scales read token-major like K's, - the nibbles swapped, the second-level scale ignored, K and V swapped -- are - all rejected by `compare`.""" - args, _ = _build(_with(segments=((2032, 272),), prefill_pair=True, global_v=2.5), seed=9000) - out = op(**args) - _compare(out, op.reference(**args)) - v_scales = args["v_scale_cache"] - pages, kv, tpb, groups = v_scales.shape - token_major = ( - v_scales.reshape(pages, kv, tpb // 4, 4, groups) - .permute(0, 1, 2, 4, 3) - .reshape(v_scales.shape) - ) - nibbles = (args["v_cache"] >> 4) | (args["v_cache"] << 4) - variants = { - "v_scales_token_major": dict(v_scale_cache=token_major), - "nibbles_swapped": dict(v_cache=nibbles), - "no_second_level_scale": dict( - kv_scale_quant_orig=torch.ones(3, dtype=torch.float32, device=DEV) - ), - "k_and_v_swapped": dict( - k_cache=args["v_cache"], - v_cache=args["k_cache"], - k_scale_cache=args["v_scale_cache"], - v_scale_cache=args["k_scale_cache"], - ), - } - for name, override in variants.items(): - wrong = op.reference(**{**args, **override}) - with pytest.raises(AssertionError): - _compare(out, wrong) - assert (out.float() - wrong.float()).abs().max().item() > 0.05, name - - def test_calls_are_bitwise_deterministic() -> None: - args, _ = _build(_with(segments=tuple((4000 + 37 * i, 2) for i in range(4))), seed=8800) + 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)) - - -_REJECTIONS = { - "float32_query": (lambda a: {**a, "q": a["q"].float()}, "q must be BF16 or FP16"), - "bf16_pool": ( - lambda a: {**a, "k_cache": a["k_cache"].view(torch.int8)}, - "k_cache must be uint8", - ), - "single_scale": ( - lambda a: {**a, "kv_scale_quant_orig": a["kv_scale_quant_orig"][:1]}, - "kv_scale_quant_orig must be a contiguous fp32 [3]", - ), - "positions_without_ratio": ( - lambda a: {**a, "query_positions": torch.zeros(4, dtype=torch.int64, device=DEV)}, - "pass both query_positions and compress_ratio, or neither", - ), - "host_query": (lambda a: {**a, "q": a["q"].cpu()}, "every operand must be on q's CUDA device"), -} - - -@pytest.mark.parametrize("case", list(_REJECTIONS), ids=list(_REJECTIONS)) -def test_rejects(case) -> None: - """The op's own loud refusals -- the reason the entry carries no `is_valid`. - One per guard family: q's dtype, the pool planes, the second-level scales, - the prefill pair and the device.""" - args, _ = _build(_with(segments=((600, 1),) * 4), seed=9100) +def test_rejects_positions_without_a_compress_ratio() -> None: + """The op's own loud refusal -- the reason the entry carries no `is_valid`: half of + the prefill pair.""" + args = _build(_with(segments=((600, 1),) * 4), seed=9100) args["selected_tokens"] = args["selected_tokens"][:, :64].contiguous() - mutate, message = _REJECTIONS[case] + args["query_positions"] = torch.zeros(4, dtype=torch.int64, device=DEV) + message = "pass both query_positions and compress_ratio, or neither" with pytest.raises(ValueError, match=re.escape(message)): with validating(op): - op(**mutate(args)) + op(**args) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_prefill_compress.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_prefill_compress.py index c1b5a3dee813..cac393b9865c 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_prefill_compress.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_prefill_compress.py @@ -90,33 +90,17 @@ def _build(spec, seed: int): table[r, :cols] = order[used : used + cols] used += cols p = torch.arange(cols * tpb).view(cols, tpb) - if spec["distinct_axes"]: - coords = (p, (3 * p + 5 * r) % 500, (5 * p + 11 * r) % 500) - else: - coords = (p, p, p) + coords = (p, (3 * p + 5 * r) % 500, (5 * p + 11 * r) % 500) position_cache[table[r, :cols].long()] = torch.stack(coords, dim=-1).to(torch.int32) assert used <= pages and len(segments) < slots, "pool or table too small for the batch" k_norm_weight = (torch.randn(head_dim, generator=gen) * 0.1).to(torch.bfloat16) logical = torch.tensor(positions, dtype=spec["logical_dtype"]) request = torch.tensor(requests, dtype=spec["request_dtype"]) - if spec["duplicate"]: - logical, request = torch.cat([logical, logical]), torch.cat([request, request]) table, index_cache, position_cache = table.to(DEV), index_cache.to(DEV), position_cache.to(DEV) extras = {} - if spec["layout"] == "strided": - wide_positions = torch.randint(0, 500, (pages, tpb, 6), dtype=torch.int32, device=DEV) - wide_positions[..., ::2] = position_cache - position_cache = wide_positions[..., ::2] - wide_table = torch.randint(0, pages, (slots, 2 * width), dtype=torch.int32, device=DEV) - wide_table[:, ::2] = table - table = wide_table[:, ::2] - pool = torch.randn(2 * pages, tpb, 1, head_dim, device=DEV).to(torch.bfloat16) - pool[pages:] = index_cache - index_cache = pool[pages:] - extras["neighbour"] = (pool[:pages], pool[:pages].clone()) - elif spec["layout"] == "shared_page": + if spec["layout"] == "shared_page": # An NVFP4 pool's page: `roles` index pages, then the position page. The other # index roles hold a sentinel and must survive, as must the position page. roles = spec["roles"] @@ -134,11 +118,8 @@ def _build(spec, seed: int): index_cache, position_cache = shared_index, shared_position extras["neighbour"] = (others, others.clone()) extras["position_page"] = (position_cache, position_cache.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()) + else: + assert spec["layout"] == "plain", spec["layout"] args = dict( logical_positions=logical.to(DEV), @@ -176,36 +157,18 @@ def test_certified_cells(cell, certifying): # noqa: F811 assert int(changed) == _anchors(args), "exactly the anchor slots are written" if "neighbour" in extras: view, saved = extras["neighbour"] - assert torch.equal(view, saved), "wrote outside the strided view" + assert torch.equal(view, saved), "wrote into another index role of the page" if "position_page" in extras: view, saved = extras["position_page"] assert torch.equal(view, saved), "wrote into the position page" def _scene(seed: int, **overrides): - """Two requests of 64 tokens from 0 under mRoPE, at the checkpoint geometry.""" + """Two requests of 64 tokens from 0, the target's call, at the checkpoint geometry.""" spec = {**op.CELLS[0].spec, "segments": ((0, 64),) * 2, "pages": 4, **overrides} return _build(spec, seed=seed)[0] -def test_text_rows_reduce_mrope_to_plain_rope() -> None: - """With t == h == w, the mRoPE selection is an identity, bit for bit.""" - args = _scene(seed=7, distinct_axes=False) - plain = dict(args, index_cache=args["index_cache"].clone(), mrope_section=None) - op(**args) - op(**plain) - assert torch.equal(args["index_cache"], plain["index_cache"]) - - -def test_a_second_call_over_the_same_rows_is_not_idempotent() -> None: - """Recorded, not tolerated: the anchor slot is one of its own inputs.""" - args = _scene(seed=14) - once = op(**args).clone() - twice = op(**args) - changed = (twice.view(torch.int16) != once.view(torch.int16)).any(-1).sum() - assert int(changed) == _anchors(args) - - #: The cell of 32 MTP7 verify blocks (256 rows, two anchors each): the segments, pool and table. _MTP7_32 = next(c.spec for c in op.CELLS if c.why.startswith("decode MTP7: 32 requests")) _MTP7_BATCH = {key: _MTP7_32[key] for key in ("segments", "pages", "slots", "rounding_flips")} @@ -238,64 +201,6 @@ def test_cuda_graph_capture_and_replay() -> None: _compare(cache, ref, batch["rounding_flips"]) -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) @@ -307,100 +212,26 @@ def _strided(tensor: torch.Tensor, dim: int) -> torch.Tensor: return wide[tuple(index)] -@pytest.mark.parametrize( - "name, dim, match", - [ - ("logical_positions", 0, "logical_positions must be contiguous"), - ("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. - - One per guard form: a 1-D operand read flat (`request_indices` and - `k_norm_weight` carry the same `is_contiguous` guard), and `cos_sin`'s last - dim. Old receipt's max abs error against the right answer: 4.72 and 7.16 - respectively; the floor below states a bound, not one draw. +def test_guard_refuses_strided_logical_positions() -> None: + """Shown unguarded first -- accepted and wrong -- then refused: a 1-D operand the + kernel reads flat (`request_indices` and `k_norm_weight` carry the same + `is_contiguous` guard). Old receipt's max abs error against the right answer: 4.72; + the floor below states a bound, not one draw. """ args = _scene(seed=12) before = args["index_cache"].clone() ref = op.reference(**args) - bad = dict(args, **{name: _strided(args[name], dim)}) - assert not bad[name].is_contiguous() + bad = dict(args, logical_positions=_strided(args["logical_positions"], 0)) + assert not bad["logical_positions"].is_contiguous() out = op(**bad) error = (out.float() - ref.float()).abs().max().item() - assert error > 0.25, f"strided {name} was expected to be read wrongly; off by {error}" + assert error > 0.25, ( + f"strided logical_positions was expected to be read wrongly; off by {error}" + ) args["index_cache"].copy_(before) - with pytest.raises(AssertionError, match=match): + with pytest.raises(AssertionError, match="logical_positions must be contiguous"): with validating(op): op(**bad) assert torch.equal(args["index_cache"], before), "a refused call still wrote" - - -def _variant(args, defect): - """The compression with one deliberate defect, in plain torch.""" - cache, cos_sin = args["index_cache"], args["cos_sin"] - cr, tpb = args["compress_ratio"], args["tokens_per_block"] - out = cache.clone() - logical = args["logical_positions"].long() - anchored = (logical + 1) % cr == 0 - anchor, request = logical[anchored], args["request_indices"].long()[anchored] - members = anchor[:, None] - (cr - 1) + torch.arange(cr, device=DEV) - pages = args["block_table"].long()[request[:, None], members // tpb] - slots = members % tpb - raw = cache[pages, slots, 0] - if defect == "pool_in_bfloat16": - acc = raw[:, 0] - for i in range(1, cr): - acc = (acc + raw[:, i]).to(torch.bfloat16) - pooled = (acc / cr).to(torch.bfloat16).float() - else: - pooled = (raw.float().sum(dim=1) / cr).to(torch.bfloat16).float() - scale = torch.rsqrt(pooled.mul(pooled).sum(-1, keepdim=True) / cache.shape[3] + args["eps"]) - normed = (pooled * scale * (args["k_norm_weight"].float() + 1.0)).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] - first, second = normed[:, :n], normed[:, n : 2 * 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), - ("ignore_mrope", 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; two of them by more than 1.0. - - Pooling in bf16 moves the result by only ~0.03 (old receipt: 2124 elements, - max 0.03125) -- inside any bf16 band, which is why the gate is bitwise. - """ - args = _scene(seed=15) - ref = op.reference(**args) - assert torch.equal(_variant(args, None).view(torch.int16), ref.view(torch.int16)) - wrong = _variant(args, defect) - out = op(**args) - _compare(out, ref) - assert int((out.view(torch.int16) != wrong.view(torch.int16)).sum()) > 0, defect - assert (out.float() - wrong.float()).abs().max().item() > floor diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py index e7eaac4ba272..fdbdc68c264d 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py @@ -19,9 +19,6 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention import ( # noqa: E501 qsa_sparse_gqa_rowtile as _module, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_sparse_gqa import ( # noqa: E501 - qsa_paged_sparse_gqa as per_row, -) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_sparse_gqa_rowtile import ( # noqa: E501 qsa_sparse_gqa_rowtile as op, ) @@ -130,8 +127,7 @@ def _build(spec, seed): rows, slots = positions.numel(), spec["slots"] assert len(segments) <= slots slot_of = torch.randperm(slots, generator=torch.Generator().manual_seed(seed))[: len(segments)] - floor = spec.get("min_seq_len", 0) - mapped = [math.ceil(max(s + n, floor) / tpb) for s, n in segments] + mapped = [math.ceil((s + n) / tpb) for s, n in segments] assert max(mapped) <= spec["pages"] pool_pages = sum(mapped) + 7 page_ids = torch.randperm(pool_pages, generator=gen, device=DEV).to(torch.int32) @@ -150,10 +146,6 @@ def _build(spec, seed): selected = _selected_tokens(positions, seed + 3, spec["compress_ratio"], spec["token_topk"], spec.get("shared_scores", False)) # fmt: skip request_indices = slot_of.to(DEV)[req].to(spec["request_dtype"]) - 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 return dict( q=q, k_cache=k_cache, @@ -180,29 +172,16 @@ def _is_prefix_packed(table): return (valid == (columns < valid.sum(1)[:, None])).all(1) -def _shared_groups(request_indices, rows_per_tile=10): - """How many 10-row groups name a single request (served by the row-tile pass).""" - rows = request_indices.numel() - padded = torch.full((math.ceil(rows / rows_per_tile) * rows_per_tile,), -1, device=DEV) - padded[:rows] = request_indices.long() - grid = padded.view(-1, rows_per_tile) - live = torch.arange(padded.numel(), device=DEV).view(-1, rows_per_tile) < rows - low = torch.where(live, grid, 1 << 40).min(1).values - high = torch.where(live, grid, -1).max(1).values - return int((low == high).sum()) - - @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 not spec.get("holes"): - # The per-row kernel's precondition, which this op shares. - selected = args["selected_tokens"] - assert bool(_is_prefix_packed(selected).all()) - assert torch.equal((selected >= 0).sum(1), _bound(args["query_positions"], spec)) + # The per-row kernel's precondition, which this op shares. + selected = args["selected_tokens"] + assert bool(_is_prefix_packed(selected).all()) + assert torch.equal((selected >= 0).sum(1), _bound(args["query_positions"], spec)) out = op(**args) assert [name for name, _ in launches] == list(_KERNELS) groups = math.ceil(rows / 10) @@ -215,10 +194,6 @@ def test_certified_cells(cell, certifying, launches) -> None: # noqa: F811 _compare(out, op.reference(**args)) -# -------------------------------------------------------------------------- -# semantics a cell cannot state -# -------------------------------------------------------------------------- - CKPT = op.CELLS[0].spec assert (CKPT["q_heads"], CKPT["kv_heads"], CKPT["head_dim"]) == (24, 2, 256) @@ -227,174 +202,14 @@ def _with(segments, **overrides): return {**CKPT, "segments": tuple(segments), **overrides} -def _plant_past_the_bound(args): - """Put each row's own newest token in its first unscanned column (bound - rounded up to 64), where the row has one.""" - positions = args["query_positions"] - final_topk = args["selected_tokens"].shape[1] - bound = _bound(positions, CKPT).clamp(max=final_topk) - first = (bound + 63) // 64 * 64 - gap_rows = (first < final_topk).nonzero().flatten() - assert gap_rows.numel() > 0 - planted = args["selected_tokens"].clone() - planted[gap_rows, first[gap_rows]] = positions[gap_rows].to(torch.int32) - return planted, gap_rows - - -@pytest.mark.parametrize( - "segments", - [ - [(0, 8191)], - [(0, 900), (907, 200), (1814, 150), (2721, 28)] - + [(100 + (499 * g) % 8000, 1) for g in range(28)], - ], - ids=["shared-groups", "mixed-with-lone-rows"], -) -def test_prefill_bound_is_honoured(segments, launches) -> None: - """An index past the per-row kernel's bound is not attended, on the - row-tile pass and on the lone-row pass alike; the planted index moves the - right answer by more than 0.1, so the comparison is not vacuous.""" - args = _build(_with(segments, slots=64), seed=11000) - planted, gap_rows = _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[gap_rows].float() - without_it[gap_rows].float()).abs().max().item() > 0.1 - assert _shared_groups(args["request_indices"]) > 0 - - -def test_repeated_index_on_a_shared_group_is_attended_once() -> None: - """The row-tile pass attends to the set a row names (the documented - precondition): a repeated column matches the reference without the - repeat. Rows of a group that spans requests (rows 10..15 here: 15 rows of - one request, then one of another) keep per-column semantics.""" - args = _build(_with([(900, 15), (5000, 1)], slots=8), seed=8700) - selected = args["selected_tokens"].clone() - once = selected.clone() - shared_row, lone_row = 3, 15 - for r in (shared_row, lone_row): - assert int(selected[r, 0]) >= 0 - last = int((selected[r] >= 0).sum()) - 1 - selected[r, last] = selected[r, 0] # the row's last token becomes a repeat of its first - once[r, last] = -1 - out = op(**{**args, "selected_tokens": selected}) - _compare(out[:10], op.reference(**{**args, "selected_tokens": once})[:10]) - _compare(out[10:], op.reference(**{**args, "selected_tokens": selected})[10:]) - - -def test_matches_the_per_row_entry_on_dumped_layout() -> None: - """Same arguments, same answer as `qsa_paged_sparse_gqa` within twice the band.""" - args = _build( - _with([(0, 1200), (3000, 900)] + [(8300 + g, 1) for g in range(9)], slots=16), seed=4242 - ) - out, theirs = op(**args), per_row(**args) - assert_within_ulp(out, theirs, element_ulp=8.0, rms_ulp=4.0) - - def test_calls_are_bitwise_deterministic() -> None: args = _build(_with([(0, 2000), (500, 300), (7000, 1)], slots=8), seed=8800) assert torch.equal(op(**args), op(**args)) -def test_zero_rows_return_an_empty_output() -> None: - args = _build(_with([(900, 4)], slots=8), seed=8400) - for name in ("q", "selected_tokens", "request_indices", "query_positions"): - args[name] = args[name][:0] - out = op(**args) - assert out.shape == (0, CKPT["q_heads"], CKPT["head_dim"]) and out.dtype == torch.bfloat16 - - -#: One loud rejection per family: the missing prefill pair, a dtype, a shape and a layout. -_REJECTIONS = { - "no_prefill_pair": ( - lambda a: {**a, "query_positions": None, "compress_ratio": None}, - "the row-tile kernel needs the prefill pair", - ), - "float32_query": (lambda a: {**a, "q": a["q"].float()}, "QSA queries must use BF16"), - "width_vs_ratio": ( - lambda a: {**a, "selected_tokens": a["selected_tokens"][:, :64].contiguous()}, - "incompatible with its compression ratio", - ), - "strided_request_indices": ( - lambda a: {**a, "request_indices": torch.zeros(8, dtype=torch.int32, device=DEV)[::2]}, - "QSA request indices must be contiguous", - ), -} - - -@pytest.mark.parametrize("case", list(_REJECTIONS), ids=list(_REJECTIONS)) -def test_rejects(case) -> None: +def test_rejects_a_call_without_the_prefill_pair() -> None: + """The op's own loud refusal: the row-tile pass is bounded by the prefill pair.""" args = _build(_with([(600, 4)], slots=8), seed=9100) - mutate, message = _REJECTIONS[case] - with pytest.raises(ValueError, match=re.escape(message)): - op(**mutate(args)) - - -class _PlanRecorder: - """Captures the plan (`group_plan`) the row-tile launch reads.""" - - def __init__(self, inner, log): - self._inner, self._log = inner, log - - def __getitem__(self, grid): - def launch(*args, **kwargs): - self._log.append(args[5].clone()) - return self._inner[grid](*args, **kwargs) - - return launch - - -def _expected_union(args, rows_per_tile=10): - """Per 10-row group: the number of distinct selected tokens within each - row's bound whose page the group's request maps, when the group names - one request (else -1), and the dense walk's tiles.""" - selected, positions = args["selected_tokens"], args["query_positions"] - request_indices, table = args["request_indices"].long(), args["block_table"] - tpb, rows = args["tokens_per_block"], selected.shape[0] - bound = _bound(positions, CKPT).clamp(max=selected.shape[1]) - columns = torch.arange(selected.shape[1], device=DEV)[None, :] - counts, dense = [], [] - for start in range(0, rows, rows_per_tile): - stop = min(start + rows_per_tile, rows) - group_requests = request_indices[start:stop].unique() - if group_requests.numel() != 1: - counts.append(-1) - dense.append(0) - continue - tokens = selected[start:stop].long() - tokens = tokens[(columns < bound[start:stop, None]) & (tokens >= 0)].unique() - pages = table[int(group_requests)][(tokens // tpb).clamp(max=table.shape[1] - 1)] - kept = tokens[(tokens // tpb < table.shape[1]) & (pages >= 0)] - counts.append(kept.numel()) - dense.append((int(tokens.max()) + 64) // 64 if tokens.numel() else 0) - return torch.tensor(counts), torch.tensor(dense) - - -@pytest.mark.parametrize( - "segments, shared_scores", - [([(4096, 4096)], True), ([(900, 40), (950, 40), (3000, 25)], False)], - ids=["sparse-union", "two-requests-unmapped-page"], -) -def test_packed_walk_is_the_resident_union(segments, shared_scores, monkeypatch) -> None: - """The row-tile pass walks exactly the union of its rows' resident - selected tokens, 64 per tile: the plan's list length is that union's - size for every one-request group, and on rows ranking one score vector - the walk is far shorter than every tile up to the group's highest token.""" - args = _build(_with(segments, slots=16, shared_scores=shared_scores), seed=7300) - if not shared_scores: - args["block_table"][:, 3] = -1 # an unmapped page inside every request's span - log: list = [] - name = "_rowtile_attention_cute" - monkeypatch.setattr(_module, name, _PlanRecorder(getattr(_module, name), log)) - out = op(**args) - _compare(out, op.reference(**args)) - plan = log[0].cpu() - counts, dense = _expected_union(args) - one_request = counts >= 0 - assert torch.equal(plan[one_request, 1], counts[one_request].int()) - taken = plan[:, 0] >= 0 - assert bool(taken.any()) - assert torch.equal(plan[taken, 0], (plan[taken, 1] + 63) // 64) - if shared_scores: - assert int(plan[taken, 0].sum()) * 2 < int(dense[taken].sum()) + args.update(query_positions=None, compress_ratio=None) + with pytest.raises(ValueError, match=re.escape("the row-tile kernel needs the prefill pair")): + op(**args) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_expand_decode.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_expand_decode.py index 84ff70dd08e9..c8e63aaa8343 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_expand_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_expand_decode.py @@ -5,10 +5,7 @@ Cells against the entry's reference (a stable descending sort, ties to the higher block, then the expansion entry's reference): the block list and the token selection exactly. Columns past each row's visible count carry garbage (huge values, NaN) that must never -be read. Beyond the cells: bitwise the composition it replaces (`qsa_topk_select`, then -`qsa_canonical_selection`, then `qsa_expand_block_indices`, each certified on its own) on -one representative decode batch and one context chunk; a stale-read check under graph -replay; one compiled kernel for every row count. +be read. Beyond the cells: a stale-read check under graph replay. """ import zlib @@ -16,21 +13,9 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_canonical_selection import ( # noqa: E501 - qsa_canonical_selection, -) -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, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_topk_expand_decode import ( # noqa: E501 - _COMPILED, -) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_topk_expand_decode import ( # noqa: E501 qsa_topk_expand_decode as op, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_topk_select import ( # noqa: E501 - qsa_topk_select, -) __extra_import_path__ = ["../../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name @@ -43,11 +28,20 @@ def _build(spec, seed): g = torch.Generator(device=DEV).manual_seed(seed) - pos = torch.tensor(spec["pos"], dtype=torch.int64, device=DEV) + if "chunk" in spec: + # One request's consecutive context rows: every row's KV length is the chunk's end. + start, count = spec["chunk"] + pos = torch.arange(start, start + count, dtype=torch.int64, device=DEV) + seqlen = torch.full((count,), start + count, dtype=torch.int32, device=DEV) + else: + pos = torch.tensor(spec["pos"], dtype=torch.int64, device=DEV) + seqlen = (pos + 1).to(torch.int32) rows = pos.numel() scores = torch.randn(rows, W, generator=g, device=DEV).abs() * 3.0 - if spec.get("ties"): - scores = torch.round(scores * 2.0) / 2.0 # few distinct values: ties at the threshold + if spec.get("relu_sparse"): + # Fewer positive scores than the budget, so the zero class straddles it. + scores = torch.relu(torch.randn(rows, W, generator=g, device=DEV) - 2.2) + scores[:, ::7] = -0.0 # negative zero ties with zero if spec.get("narrow"): scores = 1.0 + 0.05 * torch.rand(rows, W, generator=g, device=DEV) ends = ((pos + 1) // CR).to(torch.int32) @@ -56,7 +50,6 @@ def _build(spec, seed): scores[r, n:] = 1.0e30 if n + 3 < W: scores[r, n + 3] = float("nan") - seqlen = (pos + 1).to(torch.int32) topk = torch.full((rows, K), -7, dtype=torch.int32, device=DEV) return [scores, ends, pos, seqlen, topk, CR, TOKEN_TOPK] @@ -66,66 +59,25 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by certifying(op) args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) want_topk, want_sel = op.reference(*args) - sel = op(*args) + margin = cell.spec.get("out_margin") + if margin is None: + sel = op(*args) + else: + # Rows [margin, margin + rows) of a step-sized selection; the rows around stay put. + rows = args[0].shape[0] + whole = torch.full( + (rows + 2 * margin, TOKEN_TOPK + CR - 1), -9, dtype=torch.int32, device=DEV + ) + sel = op(*args, out=whole[margin : margin + rows]) + assert sel.data_ptr() == whole[margin].data_ptr() + assert (whole[:margin] == -9).all() and (whole[margin + rows :] == -9).all(), ( + "rows outside out" + ) assert sel.dtype == torch.int32 and sel.shape == want_sel.shape assert torch.equal(args[4], want_topk), "block list" assert torch.equal(sel, want_sel.to(torch.int32)), "token selection" -def _relu_sparse_rows(rows: int, seed: int) -> torch.Tensor: - """Score rows the canonical rewrite exists for: relu-sparse, so fewer positive scores - than the budget and the zero class (negative zeros among it) straddles it.""" - g = torch.Generator(device=DEV).manual_seed(seed) - scores = torch.relu(torch.randn(rows, W, generator=g, device=DEV) - 2.2) - scores[:, ::7] = -0.0 # negative zero ties with zero - return scores - - -def test_bitwise_the_canonical_composition() -> None: - """Windows across the budget (below it, at it, past it, the whole table); the block - list equals `qsa_topk_select` + `qsa_canonical_selection` and the token selection that - list's expansion, bitwise.""" - rows = 64 - scores = _relu_sparse_rows(rows, seed=zlib.crc32(b"relu_sparse") % 2**31) - g = torch.Generator(device=DEV).manual_seed(5) - pos = torch.randint(0, CR * W, (rows,), generator=g, device=DEV, dtype=torch.int64) - pos[:6] = torch.tensor([0, 4 * K - 1, 4 * K, 4 * K + 3, 8700, CR * W - 1], device=DEV) - ends = ((pos + 1) // CR).to(torch.int32) - seqlen = (pos + 1).to(torch.int32) - topk = torch.full((rows, K), -7, dtype=torch.int32, device=DEV) - sel = op(scores, ends, pos, seqlen, topk, CR, TOKEN_TOPK) - starts = torch.zeros(rows, dtype=torch.int32, device=DEV) - want = torch.full_like(topk, -7) - qsa_topk_select(scores, starts, ends, want, index_topk=K) - qsa_canonical_selection(scores, starts, ends, want) - want_sel = qsa_expand_block_indices(want, pos, seqlen, compress_ratio=CR, token_topk=TOKEN_TOPK) - assert torch.equal(topk, want), "block list" - assert torch.equal(sel, want_sel.to(torch.int32)), "token selection" - - -def test_a_context_chunk_into_a_row_range_of_the_selection() -> None: - """A context chunk's 2048 consecutive rows (one request, positions 6144..8191), written - through `out` into rows [64, 2112) of a step-sized selection: the block list and those - rows equal the canonical composition bitwise; the rows around them are untouched.""" - rows = 2048 - scores = _relu_sparse_rows(rows, seed=31) - pos = torch.arange(6144, 6144 + rows, device=DEV, dtype=torch.int64) - ends = ((pos + 1) // CR).to(torch.int32) - seqlen = torch.full((rows,), 6144 + rows, dtype=torch.int32, device=DEV) - topk = torch.full((rows, K), -7, dtype=torch.int32, device=DEV) - whole = torch.full((rows + 128, TOKEN_TOPK + CR - 1), -9, dtype=torch.int32, device=DEV) - sel = op(scores, ends, pos, seqlen, topk, CR, TOKEN_TOPK, out=whole[64 : 64 + rows]) - assert sel.data_ptr() == whole[64].data_ptr() - starts = torch.zeros(rows, dtype=torch.int32, device=DEV) - want = torch.full_like(topk, -7) - qsa_topk_select(scores, starts, ends, want, index_topk=K) - qsa_canonical_selection(scores, starts, ends, want) - want_sel = qsa_expand_block_indices(want, pos, seqlen, compress_ratio=CR, token_topk=TOKEN_TOPK) - assert torch.equal(topk, want), "block list" - assert torch.equal(whole[64 : 64 + rows], want_sel.to(torch.int32)), "token selection" - assert (whole[:64] == -9).all() and (whole[64 + rows :] == -9).all(), "rows outside out" - - def test_reads_after_the_wait_under_graph_replay() -> None: args = _build(dict(pos=(8700,)), seed=21) src = args[0].clone() @@ -149,20 +101,3 @@ def step(): got_topk = args[4].clone() sel = op(src.clone(), *args[1:]) assert torch.equal(graph_sel, sel) and torch.equal(got_topk, args[4]), f"replay {trial}" - - -def test_one_kernel_serves_every_row_count() -> None: - """The row count is a runtime argument: an eager step's context rows come in any count, - and every count here runs the one kernel compiled for the table, each matching the - reference.""" - compiled = None - for rows in (1, 2, 7, 33, 300): - g = torch.Generator(device=DEV).manual_seed(rows) - pos = torch.randint(0, CR * W, (rows,), generator=g, device=DEV, dtype=torch.int64) - args = _build(dict(pos=tuple(pos.tolist())), seed=rows) - want_topk, want_sel = op.reference(*args) - sel = op(*args) - assert torch.equal(args[4], want_topk), f"block list, {rows} rows" - assert torch.equal(sel, want_sel.to(torch.int32)), f"token selection, {rows} rows" - compiled = len(_COMPILED) if compiled is None else compiled - assert len(_COMPILED) == compiled, f"{rows} rows compiled another kernel" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_select.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_select.py index 00cdbf65da13..e0ee54787a3a 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_select.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_select.py @@ -97,42 +97,10 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 _compare(out, op.reference(*args)) -def test_selection_honours_row_ends_and_is_not_a_prefix() -> None: - """The two wrong answers that would still look plausible. - - The global maxima sit outside every window, and the in-window winners are - the *last* visible columns: a selector that ignored `row_ends` would take - the former, one that returned the first k columns would miss the latter. - """ - gen = torch.Generator(device=DEV).manual_seed(2) - rows, cols, top_k = 6, 1024, 64 - logits = torch.rand(rows, cols, generator=gen, device=DEV) - logits[:, 256:] += 100.0 - logits[:, 192:256] += 10.0 - row_starts = torch.zeros(rows, dtype=torch.int32, device=DEV) - row_ends = torch.full((rows,), 256, dtype=torch.int32, device=DEV) - indices = torch.full((rows, top_k), SENTINEL, dtype=torch.int32, device=DEV) - out = op(logits, row_starts, row_ends, indices, top_k) - _compare(out, op.reference(logits, row_starts, row_ends, indices, top_k)) - assert bool((out < 256).all()), "selected a column beyond row_ends" - assert bool((out >= 192).all()), "did not find the planted in-window maxima" - - -def test_nan_inside_the_window_is_selected() -> None: - """Why the note makes finiteness a precondition: NaN outranks every value.""" - gen = torch.Generator(device=DEV).manual_seed(11) - logits = torch.randn(1, 32, generator=gen, device=DEV) - logits[0, :4] = float("nan") - row_starts = torch.zeros(1, dtype=torch.int32, device=DEV) - row_ends = torch.full((1,), 32, dtype=torch.int32, device=DEV) - out = op(logits, row_starts, row_ends, torch.empty(1, 8, dtype=torch.int32, device=DEV), 8) - assert {0, 1, 2, 3} <= set(out[0].tolist()), "NaN is no longer selected; update the note" - - -@pytest.mark.parametrize("rows", (8, 256)) -def test_cuda_graph_capture_and_replay(rows) -> None: - """Captured, then replayed against new scores and a new causal bound: 8 rows, and the - 256 rows a verify step of 32 requests carries at max_draft_len 7.""" +def test_cuda_graph_capture_and_replay() -> None: + """Captured, then replayed against new scores and a new causal bound: the 256 rows a + verify step of 32 requests carries at max_draft_len 7.""" + rows = 256 gen = torch.Generator(device=DEV).manual_seed(8) cols = 2048 logits = _logits("relu", rows, cols, gen) @@ -159,21 +127,6 @@ def test_cuda_graph_capture_and_replay(rows) -> None: _compare(indices, op.reference(logits, row_starts, row_ends, indices, BLOCK_TOPK)) -def test_rejects_out_of_contract() -> None: - """The op's own loud refusals -- the reason `is_valid` does not restate them.""" - rows, cols, top_k = 4, 128, 8 - logits = torch.randn(rows, cols, device=DEV) - row_starts = torch.zeros(rows, dtype=torch.int32, device=DEV) - row_ends = torch.full((rows,), cols, dtype=torch.int32, device=DEV) - out = torch.empty(rows, top_k, dtype=torch.int32, device=DEV) - with pytest.raises(RuntimeError, match="expected scalar type Float"): - op(logits.to(torch.bfloat16), row_starts, row_ends, out, top_k) - with pytest.raises(RuntimeError, match="expected scalar type Int"): - op(logits, row_starts, row_ends, out.long(), top_k) - with pytest.raises(RuntimeError, match="one entry per row"): - op(logits, row_starts[:2], row_ends[:2], out, top_k) - - def _layout_case(): torch.manual_seed(10) rows, cols, top_k = 4, 128, 8 @@ -183,25 +136,6 @@ def _layout_case(): return logits, row_starts, row_ends, top_k -def test_guard_refuses_a_wider_output_buffer() -> None: - """Accepted by the op, which packs the rows at a stride of `index_topk` - regardless: row 1's answer starts inside row 0's slots and the last row is - never written. Shown unguarded first, then refused.""" - logits, row_starts, row_ends, top_k = _layout_case() - rows = logits.shape[0] - wide = torch.full((rows, top_k + 4), -99, dtype=torch.int32, device=DEV) - op(logits, row_starts, row_ends, wide, top_k) - torch.cuda.synchronize() - packed = wide.flatten()[: rows * top_k].view(rows, top_k) - _compare(packed, op.reference(logits, row_starts, row_ends, packed, top_k)) - assert not bool((wide[:, top_k:] == -99).all()), "the hazard stopped reproducing" - assert bool((wide[-1] == -99).all()), "the hazard stopped reproducing" - - with pytest.raises(AssertionError, match=r"indices must be \[num_rows, index_topk\]"): - with validating(op): - op(logits, row_starts, row_ends, wide, top_k) - - def test_guard_refuses_the_token_budget_default() -> None: """The default `index_topk` is 2048, the token budget; the target's buffer is 512 wide. That is the wider-buffer fault the other way round, and the @@ -211,14 +145,3 @@ def test_guard_refuses_the_token_budget_default() -> None: with pytest.raises(AssertionError, match="index_topk=2048"): with validating(op): op(logits, row_starts, row_ends, indices) - - -def test_guard_refuses_a_non_contiguous_output() -> None: - """A row slice of a wider arena: the same flat addressing, the same fault.""" - logits, row_starts, row_ends, top_k = _layout_case() - arena = torch.full((logits.shape[0], top_k + 4), -99, dtype=torch.int32, device=DEV) - view = arena[:, :top_k] - assert not view.is_contiguous() - with pytest.raises(AssertionError, match="contiguous"): - with validating(op): - op(logits, row_starts, row_ends, view, top_k) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qwen4_exp_qk_norm_rope.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qwen4_exp_qk_norm_rope.py index 739b3b0cc340..d5404666e2c6 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qwen4_exp_qk_norm_rope.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qwen4_exp_qk_norm_rope.py @@ -3,15 +3,11 @@ """GPU test for the qwen4_exp_qk_norm_rope entry.""" import functools -import sys import zlib import pytest import torch -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.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qwen4_exp_qk_norm_rope import ( # noqa: E501 CKPT, ) @@ -28,9 +24,6 @@ assert torch.cuda.is_available(), "qwen4_exp_qk_norm_rope requires a CUDA device" -#: The op's module: its two kernels' launchers and the token count between them. -M = sys.modules[type(op).__module__] - THETA = 1e7 EPS = 1e-6 TABLE_ROWS = 262144 # the checkpoint's max_position_embeddings @@ -63,7 +56,7 @@ def _distinct_rows(tokens: int, max_position: int, g: torch.Generator) -> torch. return rows.to(torch.int32).contiguous() -def _build(spec, seed, positions=None): +def _build(spec, seed): """`(proj, scratch, kwargs)`: the fused projection row, the index scratch whose first `index_heads` heads are the `index_q` view, and the op's keyword arguments. Norm weights near 1 put q/k at their largest magnitude @@ -78,8 +71,7 @@ def _build(spec, seed, positions=None): def weight(n): return (torch.randn(n, generator=g, device="cuda") * 0.1 + 1.0).to(torch.bfloat16) - if positions is None: - positions = _distinct_rows(t, spec.get("max_position", 32768), g) + positions = _distinct_rows(t, spec.get("max_position", 32768), g) kwargs = {k: v for k, v in spec.items() if k not in _NOT_OP_ARGS} kwargs.update( index_q=scratch[:, : spec["index_heads"]], @@ -136,73 +128,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 ) -def _current_pair(proj, scratch, kwargs, index_k_weight): - """What the forward ran before this entry: the CUDA norm+rope in place on - the projection's main heads, the index columns copied into the scratch, and - the CUDA norm+rope again on the scratch (row 0, plain rope).""" - spec = CKPT - tail = spec["proj_width"] // spec["head_dim"] - spec["num_heads_q"] - spec["num_heads_k"] - pos3 = kwargs["position_ids"] - fused_qk_norm_rope( - proj, num_heads_q=spec["num_heads_q"], num_heads_k=spec["num_heads_k"], - num_heads_v=tail, head_dim=spec["head_dim"], rotary_dim=spec["rotary_dim"], eps=EPS, - q_weight=kwargs["q_weight"], k_weight=kwargs["k_weight"], base=THETA, is_neox=True, - position_ids=pos3, is_qk_norm=True, use_gemma=True, use_mrope=True, - mrope_section1=spec["mrope_section1"], mrope_section2=spec["mrope_section2"], - ) # fmt: skip - flat = scratch.reshape(proj.shape[0], -1) - flat.copy_(proj[:, spec["index_q_col"] :]) - fused_qk_norm_rope( - flat, num_heads_q=spec["index_heads"], num_heads_k=1, num_heads_v=1, - head_dim=spec["index_head_dim"], rotary_dim=spec["rotary_dim"], eps=EPS, - q_weight=kwargs["index_q_weight"], k_weight=index_k_weight, base=THETA, is_neox=True, - position_ids=pos3[0].contiguous(), is_qk_norm=True, use_gemma=True, use_mrope=False, - mrope_section1=0, mrope_section2=0, - ) # fmt: skip - return proj, scratch[:, : spec["index_heads"]] - - -@pytest.mark.parametrize("tokens, rows", [(1, "distinct"), (128, "text")]) -def test_matches_the_cuda_pair_it_replaces(tokens: int, rows: str) -> None: - """Same wiring as the two CUDA calls: mRoPE sections (11, 10) on the main - plane, coordinate 0 alone on the index plane, q/k weight assignment. Both - sides are ~1 ulp from fp64 (the CUDA op's `__sincosf` phase error grows - with position), so the band is the sum of the two entries' element bands; - a wrong wiring lands ~100x outside it.""" - spec = dict(tokens=tokens, **CKPT) - g = torch.Generator(device="cuda").manual_seed(_seed(f"pair{tokens}{rows}")) - if rows == "distinct": - positions = _distinct_rows(tokens, 9216, g) - else: - row = torch.randint(8192, 9216, (tokens,), generator=g, device="cuda", dtype=torch.int32) - positions = torch.stack([row, row, row]).contiguous() - proj, scratch, kwargs = _build(spec, seed=_seed(f"pair{tokens}{rows}"), positions=positions) - old_proj, old_scratch = proj.clone(), scratch.clone() - index_k_weight = torch.ones(CKPT["index_head_dim"], dtype=torch.bfloat16, device="cuda") - old_proj, old_iq = _current_pair(old_proj, old_scratch, kwargs, index_k_weight) - new_proj, new_iq = op(proj, **kwargs) - main = _main_cols(spec) - print(f"\n main {_deviation(new_proj[:, :main], old_proj[:, :main])}") - print(f" index {_deviation(new_iq.reshape(tokens, -1), old_iq.reshape(tokens, -1))}") - assert_within_ulp(new_proj[:, :main], old_proj[:, :main], element_ulp=3.0, rms_ulp=1.0) - assert_within_ulp( - new_iq.reshape(tokens, -1), old_iq.reshape(tokens, -1), element_ulp=3.0, rms_ulp=1.0 - ) - - -def test_index_plane_sections_select_coordinates() -> None: - """Non-zero index sections read coordinates 1 and 2: the result moves away - from the plain (0, 0) one when the rows differ.""" - spec = dict(tokens=8, **CKPT) - proj, scratch, kwargs = _build(spec, seed=_seed("index-sections")) - plain = op.reference(proj, **kwargs)[1] - kwargs.update(index_mrope_section1=11, index_mrope_section2=10) - ref = op.reference(proj, **kwargs)[1] - out = op(proj.clone(), **kwargs)[1] - assert (ref.float() - plain.float()).abs().max() > 1.0 - assert_within_ulp(out.reshape(8, -1), ref.reshape(8, -1), element_ulp=2.0, rms_ulp=0.25) - - def test_cuda_graph_replay_reads_fresh_inputs() -> None: """Captured once at M=4, replayed after the inputs are rewritten in place: the replay equals an eager call on the new inputs, bit for bit.""" @@ -229,133 +154,9 @@ def test_cuda_graph_replay_reads_fresh_inputs() -> None: assert torch.equal(kwargs["index_q"], eager_kwargs["index_q"]) -def test_zero_tokens_launches_nothing() -> None: - proj, scratch, kwargs = _build(dict(tokens=0, **CKPT), seed=1, positions=torch.empty( - 3, 0, dtype=torch.int32, device="cuda")) # fmt: skip - out_proj, out_iq = op(proj, **kwargs) - assert out_proj.shape[0] == 0 and out_iq.shape[0] == 0 - - -@pytest.mark.parametrize( - "mutate, message", - [ - (lambda kw: kw.update(index_q_col=1024), "overlap"), - (lambda kw: kw["position_ids"].__setitem__((0, 0), TABLE_ROWS), "past the rope table"), - (lambda kw: kw.update(position_ids=kw["position_ids"].t().contiguous().t()), "contiguous"), - (lambda kw: kw.update(index_q=kw["index_q"][:, :3]), ""), - ], - ids=["planes-overlap", "position-past-table", "strided-positions", "index-q-heads"], -) -def test_guard_refuses(mutate, message) -> None: +def test_guard_refuses_a_position_past_the_table() -> None: + """An out-of-range position would read past the rope table.""" proj, _, kwargs = _build(dict(tokens=3, **CKPT), seed=_seed("guard")) - mutate(kwargs) - with validating(op), pytest.raises(AssertionError, match=message): + kwargs["position_ids"][0, 0] = TABLE_ROWS + with validating(op), pytest.raises(AssertionError, match="past the rope table"): op(proj, **kwargs) - - -# ---- the CuTe kernel (prefill-size token counts) against the Triton kernel ---------------- - - -def _sections(kwargs): - return ( - kwargs["mrope_section1"], - kwargs["mrope_section2"], - kwargs.get("index_mrope_section1", 0), - kwargs.get("index_mrope_section2", 0), - ) - - -def _triton(proj, kwargs): - """The Triton kernel at any token count (from 65 tokens its eight-head launch).""" - M._launch_triton( - proj, kwargs["index_q"], kwargs["position_ids"], kwargs["cos_sin"], kwargs["q_weight"], - kwargs["k_weight"], kwargs["index_q_weight"], kwargs["num_heads_q"], - kwargs["num_heads_k"], kwargs["head_dim"], kwargs["index_heads"], - kwargs["index_head_dim"], kwargs["index_q_col"], kwargs["rotary_dim"], kwargs["eps"], - _sections(kwargs), - ) # fmt: skip - - -def _with_signed_zeros(proj) -> None: - """Rows whose sums and outputs turn on zero signs: +-0 and a tiny value in a rotary chunk, - -0 across a rotary/plain boundary, a whole -0 head, a zero index head.""" - proj[0, :8] = torch.tensor([0.0, -0.0, 1e-30, -0.0, 3.0, -2.0, 0.0, -0.0]) - proj[1, 300:308] = -0.0 - proj[2, 3 * 256 : 4 * 256] = -0.0 - proj[3, CKPT["index_q_col"] : CKPT["index_q_col"] + 128] = 0.0 - - -def _bitwise(a, b) -> bool: - return torch.equal(a.view(torch.int16), b.view(torch.int16)) - - -@pytest.mark.parametrize("tokens, rows", [(1024, "distinct"), (8193, "text"), (16384, "distinct")]) -def test_cute_kernel_is_the_triton_kernel_bitwise(tokens: int, rows: str) -> None: - """From `_CUTE_MIN_TOKENS` tokens the op runs the CuTe kernel: every output bit (both - planes, the untouched columns and scratch heads included) is the Triton kernel's.""" - assert tokens >= M._CUTE_MIN_TOKENS - spec = dict(tokens=tokens, **CKPT) - proj, scratch, kwargs = _build(spec, seed=_seed(f"cute{tokens}{rows}")) - if rows == "text": - kwargs["position_ids"] = kwargs["position_ids"][:1].expand(3, tokens).contiguous() - _with_signed_zeros(proj) - ref_proj, ref_scratch = proj.clone(), scratch.clone() - _triton(ref_proj, dict(kwargs, index_q=ref_scratch[:, : CKPT["index_heads"]])) - op(proj, **kwargs) - assert _bitwise(proj, ref_proj), "projection" - assert _bitwise(scratch, ref_scratch), "index scratch" - - -def test_cute_kernel_index_sections_bitwise() -> None: - """Non-zero index sections on the CuTe kernel: the index plane reads coordinates 1 and 2 - as the Triton kernel does.""" - spec = dict(tokens=2048, **CKPT) - proj, scratch, kwargs = _build(spec, seed=_seed("cute-index-sections")) - kwargs.update(index_mrope_section1=11, index_mrope_section2=10) - ref_proj, ref_scratch = proj.clone(), scratch.clone() - _triton(ref_proj, dict(kwargs, index_q=ref_scratch[:, : CKPT["index_heads"]])) - op(proj, **kwargs) - assert _bitwise(proj, ref_proj) and _bitwise(scratch, ref_scratch) - - -def test_cute_kernel_graph_replay_reads_fresh_inputs() -> None: - """The CuTe kernel captured once and replayed after the inputs are rewritten in place.""" - spec = dict(tokens=2048, **CKPT) - proj, scratch, kwargs = _build(spec, seed=_seed("cute-graph")) - proj_new, _, kwargs_new = _build(spec, seed=_seed("cute-graph-new")) - eager_proj, eager_scratch = proj_new.clone(), scratch.clone() - eager_kwargs = dict(kwargs_new, index_q=eager_scratch[:, : CKPT["index_heads"]]) - op(eager_proj, **eager_kwargs) - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - op(proj.clone(), **dict(kwargs, index_q=scratch.clone()[:, : CKPT["index_heads"]])) - torch.cuda.current_stream().wait_stream(stream) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - op(proj, **kwargs) - proj.copy_(proj_new) - for name in ("position_ids", "q_weight", "k_weight", "index_q_weight"): - kwargs[name].copy_(kwargs_new[name]) - graph.replay() - torch.cuda.synchronize() - assert torch.equal(proj, eager_proj) - assert torch.equal(kwargs["index_q"], eager_kwargs["index_q"]) - - -def test_cute_kernel_declines_strided_index_heads(monkeypatch) -> None: - """Index heads 136 elements apart are not the CuTe kernel's layout: the Triton kernel runs.""" - tokens = 2048 - spec = dict(tokens=tokens, **CKPT) - proj, _, kwargs = _build(spec, seed=_seed("cute-declines")) - wide = torch.zeros(tokens, CKPT["index_heads"], 136, dtype=torch.bfloat16, device="cuda") - kwargs["index_q"] = wide[:, :, : CKPT["index_head_dim"]] - ref_proj, ref_wide = proj.clone(), wide.clone() - _triton(ref_proj, dict(kwargs, index_q=ref_wide[:, :, : CKPT["index_head_dim"]])) - - def refuse(*args, **kwargs): - raise AssertionError("the CuTe kernel ran on strided index heads") - - monkeypatch.setattr(M, "_launch_cute", refuse) - op(proj, **kwargs) - assert _bitwise(proj, ref_proj) and _bitwise(wide, ref_wide) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_vision_rope_2d.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_vision_rope_2d.py index 8843c837443b..2a83d03b3fc4 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_vision_rope_2d.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_vision_rope_2d.py @@ -4,12 +4,11 @@ `op.reference` shares no code with the entry: positions from nested loops, angles in fp64 from `theta ** (-2j / rope_dim)`, rotation as a complex -multiply. Four ways of getting a 2-D vision rope subtly wrong are run too, so -the gate reads against the distance to a wrong answer. +multiply. """ import zlib -from typing import List, Tuple +from typing import Tuple import pytest import torch @@ -98,146 +97,8 @@ def test_certified_cells(cell, certifying): # noqa: F811 _compare((q_out, k_out), (q_ref, k_ref)) -# --- test-local positions and angles, for the structural checks and the wrong -# operators below ------------------------------------------------------------ - - -def _positions(grid, merge) -> List[Tuple[int, int]]: - positions: List[Tuple[int, int]] = [] - for frames, height, width in grid: - frame = [] - for block_row in range(height // merge): - for block_col in range(width // merge): - for intra_row in range(merge): - for intra_col in range(merge): - frame.append((block_row * merge + intra_row, block_col * merge + intra_col)) - positions.extend(frame * frames) - return positions - - -def _angles(positions, rope_dim, theta, inv_scale_dim=None) -> torch.Tensor: - """`[tokens, rope_dim]` fp32 on cuda; height half first, then width.""" - axis_dim = rope_dim // 2 - denom = rope_dim if inv_scale_dim is None else inv_scale_dim - inverse = torch.tensor( - [theta ** (-(2.0 * j) / denom) for j in range(axis_dim)], dtype=torch.float64 - ) - pos = torch.tensor(positions, dtype=torch.float64).reshape(-1, 2) - return (pos[:, :, None] * inverse).flatten(1).to(device="cuda", dtype=torch.float32) - - -def _rotate_by_pairs(x, angles) -> torch.Tensor: - rope_dim = angles.shape[-1] - fp32 = x.float() - pair = torch.complex(fp32[..., :rope_dim], fp32[..., rope_dim:]) - rotated = pair * torch.polar(torch.ones_like(angles), angles).unsqueeze(-2) - return torch.cat((rotated.real, rotated.imag), dim=-1).to(x.dtype) - - -def _wrong_interleaved(q, grid): - """GPT-J pairing: dims (2j, 2j+1) instead of (j, j + rope_dim).""" - rope_dim = q.shape[-1] // 2 - angles = _angles(_positions(grid, MERGE), rope_dim, THETA) - pair = torch.view_as_complex(q.float().unflatten(-1, (rope_dim, 2)).contiguous()) - rotor = torch.polar(torch.ones_like(angles), angles).unsqueeze(-2) - return torch.view_as_real(pair * rotor).flatten(-2).to(q.dtype) - - -def _wrong_axes_swapped(q, grid): - rope_dim = q.shape[-1] // 2 - axis_dim = rope_dim // 2 - angles = _angles(_positions(grid, MERGE), rope_dim, THETA) - return _rotate_by_pairs(q, torch.cat((angles[:, axis_dim:], angles[:, :axis_dim]), dim=-1)) - - -def _wrong_row_major(q, grid): - positions = [] - for frames, height, width in grid: - positions.extend([(r, c) for r in range(height) for c in range(width)] * frames) - return _rotate_by_pairs(q, _angles(positions, q.shape[-1] // 2, THETA)) - - -def _wrong_full_head_inv_freq(q, grid): - head_dim = q.shape[-1] - angles = _angles(_positions(grid, MERGE), head_dim // 2, THETA, inv_scale_dim=head_dim) - return _rotate_by_pairs(q, angles) - - -WRONG_VARIANTS = { - "interleaved_pairing": _wrong_interleaved, - "axes_swapped": _wrong_axes_swapped, - "row_major_positions": _wrong_row_major, - "full_head_inv_freq": _wrong_full_head_inv_freq, -} - - -@pytest.mark.parametrize("variant", sorted(WRONG_VARIANTS)) -def test_discrimination_against_wrong_operators(variant: str) -> None: - """At 1024 bf16 tokens each wrong operator lands > 0.5 away (measured - 8.5-9.5) and > 100x the correct deviation; the entry's gate refuses it.""" - grid = [(1, 32, 32)] - q, k = _operands(_num_tokens(grid), HEADS, HEADS, HEAD_DIM, torch.bfloat16, seed=23) - grid_thw = torch.tensor(grid, dtype=torch.int32) - - q_out, k_out = op(q, k, grid_thw, MERGE, THETA) - q_ref, k_ref = op.reference(q, k, grid_thw, MERGE, THETA) - correct = _max_deviation(q_out, q_ref) - wrong_q = WRONG_VARIANTS[variant](q, grid) - wrong = _max_deviation(wrong_q, q_ref) - print(f"[discrimination/{variant}] correct={correct:.3e} wrong={wrong:.3e}") - assert wrong > 0.5, f"{variant} is only {wrong} away -- not a discriminating case" - assert wrong > 100.0 * max(correct, 1e-6) - with pytest.raises(AssertionError): - _compare((wrong_q, k_out), (q_ref, k_ref)) - - -def test_origin_token_is_the_identity() -> None: - """Patch (0, 0) has every angle zero: bitwise unchanged. No other is.""" - grid = [(1, 4, 4)] - q, k = _operands(16, HEADS, HEADS, HEAD_DIM, torch.bfloat16, seed=3) - q_out, k_out = op(q, k, torch.tensor(grid, dtype=torch.int32), MERGE, THETA) - assert torch.equal(q_out[0], q[0]) and torch.equal(k_out[0], k[0]) - assert not torch.equal(q_out[1:], q[1:]) - - -def test_rotation_preserves_the_pair_norms() -> None: - grid = [(1, 16, 24)] - q, k = _operands(_num_tokens(grid), HEADS, HEADS, HEAD_DIM, torch.float32, seed=11) - rope_dim = HEAD_DIM // 2 - q_out, k_out = op(q, k, torch.tensor(grid, dtype=torch.int32), MERGE, THETA) - for before, after in ((q, q_out), (k, k_out)): - pair_before = before[..., :rope_dim] ** 2 + before[..., rope_dim:] ** 2 - pair_after = after[..., :rope_dim] ** 2 + after[..., rope_dim:] ** 2 - torch.testing.assert_close(pair_after, pair_before, rtol=1e-5, atol=1e-5) - - -def test_height_and_width_occupy_separate_halves_of_the_angle_table() -> None: - """A column-0 token moves only in the height lanes of both rotation halves, - a row-0 token only in the width lanes.""" - grid = [(1, 4, 4)] - q, k = _operands(16, HEADS, HEADS, HEAD_DIM, torch.float32, seed=17) - rope_dim = HEAD_DIM // 2 - axis_dim = rope_dim // 2 - q_out, _ = op(q, k, torch.tensor(grid, dtype=torch.int32), MERGE, THETA) - positions = _positions(grid, MERGE) - row_only = [i for i, (r, c) in enumerate(positions) if r != 0 and c == 0] - col_only = [i for i, (r, c) in enumerate(positions) if r == 0 and c != 0] - assert row_only and col_only - for t in row_only: - torch.testing.assert_close(q_out[t, :, axis_dim:rope_dim], q[t, :, axis_dim:rope_dim]) - torch.testing.assert_close( - q_out[t, :, rope_dim + axis_dim :], q[t, :, rope_dim + axis_dim :] - ) - assert not torch.equal(q_out[t, :, :axis_dim], q[t, :, :axis_dim]) - for t in col_only: - torch.testing.assert_close(q_out[t, :, :axis_dim], q[t, :, :axis_dim]) - torch.testing.assert_close( - q_out[t, :, rope_dim : rope_dim + axis_dim], q[t, :, rope_dim : rope_dim + axis_dim] - ) - assert not torch.equal(q_out[t, :, axis_dim:rope_dim], q[t, :, axis_dim:rope_dim]) - - def test_grid_placement_and_integer_dtype_do_not_change_the_result() -> None: + """The target hands over an int64 grid on the device; the cells pass host grids.""" grid = [(1, 16, 24), (1, 8, 12)] q, k = _operands(_num_tokens(grid), HEADS, HEADS, HEAD_DIM, torch.bfloat16, seed=29) host = torch.tensor(grid, dtype=torch.int32) @@ -247,63 +108,6 @@ def test_grid_placement_and_integer_dtype_do_not_change_the_result() -> None: assert torch.equal(q_host, q_other) and torch.equal(k_host, k_other) -def test_theta_changes_the_result() -> None: - grid = [(1, 4, 4)] - q, k = _operands(16, HEADS, HEADS, HEAD_DIM, torch.float32, seed=31) - grid_thw = torch.tensor(grid, dtype=torch.int32) - default, _ = op(q, k, grid_thw, MERGE, THETA) - other, _ = op(q, k, grid_thw, MERGE, 1.0e6) - assert not torch.equal(default, other) - - -def test_cuda_graph_replay_matches_eager() -> None: - """Capturable; the grid is baked in at capture.""" - grid_thw = torch.tensor([(1, 16, 16)], dtype=torch.int32) - q, k = _operands(256, HEADS, HEADS, HEAD_DIM, torch.bfloat16, seed=37) - eager_q, eager_k = op(q, k, grid_thw, MERGE, THETA) - - warmup = torch.cuda.Stream() - warmup.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(warmup): - for _ in range(3): - op(q, k, grid_thw, MERGE, THETA) - torch.cuda.current_stream().wait_stream(warmup) - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - captured_q, captured_k = op(q, k, grid_thw, MERGE, THETA) - graph.replay() - torch.cuda.synchronize() - assert torch.equal(captured_q, eager_q) and torch.equal(captured_k, eager_k) - - -# --- guards ------------------------------------------------------------------ -# -# The indivisible-side and short-grid hazards are not driven unguarded here: -# the unwritten position rows are uninitialized memory, and an out-of-range -# value there trips a device-side assert in the angle gather that poisons the -# CUDA context for every later test in the process. - - -def test_guard_refuses_a_grid_side_not_divisible_by_the_merge() -> None: - grid = [(1, 4, 5)] - q, k = _operands(20, HEADS, HEADS, HEAD_DIM, torch.bfloat16, seed=2) - with pytest.raises(AssertionError, match="divisible by spatial_merge_size"): - with validating(op): - op(q, k, torch.tensor(grid, dtype=torch.int32), MERGE, THETA) - - -def test_guard_refuses_a_token_count_that_does_not_match_the_grid() -> None: - q, k = _operands(17, HEADS, HEADS, HEAD_DIM, torch.bfloat16, seed=2) - with pytest.raises(AssertionError, match="describes 16 tokens but q/k carry 17"): - with validating(op): - op(q, k, torch.tensor([(1, 4, 4)], dtype=torch.int32), MERGE, THETA) - q, k = _operands(64, HEADS, HEADS, HEAD_DIM, torch.bfloat16, seed=2) - with pytest.raises(AssertionError, match="describes 1 tokens but q/k carry 64"): - with validating(op): - op(q, k, torch.tensor([(1, 1, 1)], dtype=torch.int32), 1, THETA) - - def test_guard_refuses_a_one_token_k() -> None: """The hazard, unguarded: a [1, H, D] k broadcasts to [T, H, D].""" grid_thw = torch.tensor([(1, 4, 4)], dtype=torch.int32) @@ -314,13 +118,3 @@ def test_guard_refuses_a_one_token_k() -> None: with pytest.raises(AssertionError, match="must agree on tokens and head_dim"): with validating(op): op(q, k[:1], grid_thw, MERGE, THETA) - - -def test_head_dim_not_divisible_by_4_is_loud() -> None: - """Not guarded: the angle table comes out wider than the head and the - broadcast raises.""" - grid_thw = torch.tensor([(1, 4, 4)], dtype=torch.int32) - for head_dim in (70, 71): - q, k = _operands(16, HEADS, HEADS, head_dim, torch.bfloat16, seed=7) - with pytest.raises(RuntimeError): - op(q, k, grid_thw, MERGE, THETA) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/embedding/test_modeling_v2_ngram_hash.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/embedding/test_modeling_v2_ngram_hash.py index 91871a7408c3..a402400fc12d 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/embedding/test_modeling_v2_ngram_hash.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/embedding/test_modeling_v2_ngram_hash.py @@ -3,9 +3,8 @@ """GPU test for the ngram_hash entry. Exact integer arithmetic, gated bitwise against the entry's scalar reference. -A wrong hash is still a valid row index, so beyond the cells: six near-miss -derivations must each differ, the per-head bands, the EOS reset semantics, -the tail slice and chunked prefill, and the four guards. +Beyond the cells: the EOS reset that `ngram_token_history`'s EOS seed relies +on, the narrow-multiplier guard and graph capture. """ import zlib @@ -41,7 +40,6 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: DEV = "cuda" EOS = CKPT_EOS_TOKEN_ID -NGRAM_SIZE = len(CKPT_MULTIPLIERS) NGRAM_HEADS = len(CKPT_HEAD_VOCAB_SIZES) PADDED_TABLE_ROWS = 320001536 # sum(head_vocab_sizes) rounded up to a multiple of 128 @@ -91,74 +89,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res _compare(out, op.reference(*args)) -def _variant(rows, num_query_tokens, variant): - """The scalar formula with one specific mistake in it.""" - multipliers = CKPT_MULTIPLIERS[::-1] if variant == "multipliers_reversed" else CKPT_MULTIPLIERS - heads_per_ngram = NGRAM_HEADS // (NGRAM_SIZE - 1) - out = [] - for row in rows: - segment_start, last_eos = [], -1 - for position, token in enumerate(row): - if variant == "inclusive_segment" and token == EOS: - last_eos = position - segment_start.append(last_eos + 1) - if variant != "inclusive_segment" and token == EOS: - last_eos = position - row_out = [] - for position in range(len(row) - num_query_tokens, len(row)): - context = [] - for shift in range(NGRAM_SIZE): - source = position - shift - (1 if variant == "extra_shift" and shift else 0) - if variant == "no_eos_reset": - inside = source >= 0 - else: - inside = position - segment_start[position] >= shift and source >= 0 - context.append(row[source] if inside else EOS) - heads = [] - for ngram in range(2, NGRAM_SIZE + 1): - mixed = context[0] * multipliers[0] - for index in range(1, ngram): - mixed ^= context[index] * multipliers[index] - block = ngram - 2 - if variant == "blocks_swapped": - block = NGRAM_SIZE - 2 - block - low = block * heads_per_ngram - for head in range(heads_per_ngram): - offset = 0 if variant == "no_offsets" else CKPT_HEAD_OFFSETS[low + head] - heads.append(mixed % CKPT_HEAD_VOCAB_SIZES[low + head] + offset) - row_out.append(heads) - out.append(row_out) - return torch.tensor(out, dtype=torch.long, device=DEV) - - -def test_wrong_variants_are_separated() -> None: - """Six near misses: each in range, inside the table, and different.""" - ids = _tokens(2, CKPT_CONTEXT_LEN + 512, seed=5, eos_positions=(30, 31, 200, 400)) - got = _hash(ids, 512) - _compare(got, op.reference(ids, *_constants(), EOS, 512)) - for variant in ( - "extra_shift", - "no_eos_reset", - "inclusive_segment", - "multipliers_reversed", - "no_offsets", - "blocks_swapped", - ): - wrong = _variant(ids.tolist(), 512, variant) - differing = int((got != wrong).sum()) - print(f" {variant:<22} {differing} indices differ") - assert differing > 0, f"{variant} is indistinguishable from the correct result" - assert 0 <= int(wrong.min()) and int(wrong.max()) < PADDED_TABLE_ROWS - - -def test_every_head_lands_in_its_own_band() -> None: - got = _hash(_tokens(2, CKPT_CONTEXT_LEN + 512, seed=9, eos_positions=(100,)), 512) - for head in range(NGRAM_HEADS): - column = got[:, :, head] - assert int(column.min()) >= CKPT_HEAD_OFFSETS[head] - assert int(column.max()) < CKPT_HEAD_OFFSETS[head] + CKPT_HEAD_VOCAB_SIZES[head] - - def test_eos_restarts_the_window() -> None: """One past an EOS hashes as a sequence's first token, and only there.""" token = 4242 @@ -169,67 +99,6 @@ def test_eos_restarts_the_window() -> None: assert not torch.equal(_hash(fresh, 1), _hash(no_eos, 1)) -def test_eos_token_itself_ends_its_segment() -> None: - """An EOS hashes with the segment it closes, not as a fresh start.""" - closing = torch.tensor([[11, 22, 33, EOS]], device=DEV) - opening = torch.tensor([[EOS] * (CKPT_CONTEXT_LEN + 1)], device=DEV) - assert not torch.equal(_hash(closing, 1), _hash(opening, 1)) - - -def test_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() @@ -242,13 +111,6 @@ def test_guard_refuses_narrow_multipliers() -> None: op(ids, narrow, sizes, offsets, EOS, 4) -def test_single_multiplier_raises_on_its_own() -> None: - ids = _tokens(1, CKPT_CONTEXT_LEN + 4, seed=43) - multipliers, sizes, offsets = _constants() - with pytest.raises(ZeroDivisionError): - op(ids, multipliers[:1], sizes, offsets, EOS, 4) - - def test_cuda_graph_capture_and_replay() -> None: ids = _tokens(8, CKPT_CONTEXT_LEN + 1, seed=53, lead_eos=False, eos_density=0.2) constants = _constants() diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_bf16_gemv_pdl.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_bf16_gemv_pdl.py index bc27c7bfed48..ba0d3a05c9df 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_bf16_gemv_pdl.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_bf16_gemv_pdl.py @@ -5,9 +5,8 @@ Cells against the fp32 reference within the band `low_m_bf16_mm` uses (the same products, summed in another order). Beyond the cells: a PDL chain under CUDA-graph capture and replay, where each GEMV consumes the previous one's -output and so would read a stale activation if it read before its wait; PDL on -and off agree bitwise; the band rejects the near misses of a wrong layout; and -the guards. +output and so would read a stale activation if it read before its wait; and +the layout guard. """ import zlib @@ -15,10 +14,6 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.gemm.bf16_gemv_pdl import ( # noqa: E501 - MAX_ROWS, - launch, -) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.gemm.bf16_gemv_pdl import ( # noqa: E501 bf16_gemv_pdl as op, ) @@ -97,33 +92,7 @@ def chain(): h = out -@pytest.mark.parametrize("m,n,k", [(1, 16480, 2560), (1, 2560, 6144)]) -def test_pdl_on_and_off_agree_bitwise(m, n, k) -> None: - """PDL changes when the kernel may start, not what it computes.""" - mat_a, mat_b = _build(m, n, k, seed=11) - assert torch.equal(launch(mat_a, mat_b, True), launch(mat_a, mat_b, False)) - - -def test_band_discriminates_wrong_variants() -> None: - """The band admits the kernel and rejects the near misses a wrong k-tile - split or layout would produce: half of K missing, a one-column shift of the - output, a 3% scale error.""" - m, n, k = 1, 16480, 2560 - mat_a, mat_b = _build(m, n, k, seed=7) - out = op(mat_a, mat_b) - ref = op.reference(mat_a, mat_b) - _compare(out, ref) - half = mat_a.clone() - half[:, k // 2 :] = 0 - for wrong in (op.reference(half, mat_b), torch.roll(ref, 1, dims=1), ref * 1.03): - with pytest.raises(AssertionError): - _compare(out, wrong) - - -def test_guard_refuses_a_row_major_mat_b_and_too_many_rows() -> None: +def test_guard_refuses_a_row_major_mat_b() -> None: mat_a, mat_b = _build(1, 512, 2560, seed=3) with validating(op), pytest.raises(AssertionError): op(mat_a, mat_b.contiguous()) - a, b = _build(MAX_ROWS + 1, 512, 2560, seed=4) - with validating(op), pytest.raises(AssertionError): - op(a, b) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_low_m_bf16_mm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_low_m_bf16_mm.py index 303ea44b6ecd..d3b3a17e22a0 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_low_m_bf16_mm.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_low_m_bf16_mm.py @@ -6,8 +6,8 @@ branch of the table's other shapes and M buckets, so a table edit that moves a shape to another kernel fails here before it can change a number silently. Beyond the cells: CUDA-graph capture and replay after warmup (the way the -decode forward runs it), the wrong-variant discrimination of the band, and the -guards. +decode forward runs it), and the guard against a padded `mat_b`, which the +cuBLAS branch would read as dense. """ import zlib @@ -66,6 +66,9 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 _PLAN_PINS = ( (1, 10240, 2560, "direct"), (1, 48, 2560, "direct"), + (32, 48, 2560, "direct"), + (8, 2560, 6144, "splitk"), + (2, 16480, 2560, "cublas"), (1, 2560, 6144, "direct"), (3, 12288, 2560, "splitk"), (8, 12288, 2560, "splitk"), @@ -124,22 +127,6 @@ def test_graph_replay_after_warmup() -> None: _compare(out, op.reference(a, b)) -def test_band_discriminates_wrong_variants() -> None: - """The band admits the kernel and rejects the near misses a wrong layout or - a dropped K-split would produce: half of K missing, a one-column shift of - the output, a 3% scale error.""" - m, n, k = 8, 2560, 6144 - mat_a, mat_b = _build(m, n, k, seed=7) - out = op(mat_a, mat_b) - ref = op.reference(mat_a, mat_b) - _compare(out, ref) - half = mat_a.clone() - half[:, k // 2 :] = 0 - for wrong in (op.reference(half, mat_b), torch.roll(ref, 1, dims=1), ref * 1.03): - with pytest.raises(AssertionError): - _compare(out, wrong) - - def _padded(m, n, k): """A row-padded mat_a and a column-padded mat_b: dense along K, leading dimension 2K. Both are views a caller gets by slicing a wider tensor.""" @@ -159,35 +146,3 @@ def test_guard_refuses_a_padded_mat_b() -> None: with pytest.raises(AssertionError, match="column-major"): with validating(op): op(a, b) - - -def test_guard_refuses_a_padded_mat_a() -> None: - """Same for mat_a's row stride on the cuBLAS branch.""" - a, b = _padded(4, 640, 640) - b = b.t().contiguous().t() - assert plan(4, 640, 640).kind == "cublas" and a.stride() == (1280, 1) - with pytest.raises(AssertionError): - _compare(op(a, b), op.reference(a, b)) - with pytest.raises(AssertionError, match="row-major"): - with validating(op): - op(a, b) - - -def test_flashinfer_branches_reject_padded_operands_loudly() -> None: - """The FlashInfer branches validate layouts themselves and raise.""" - a, b = _padded(1, 10240, 2560) - assert plan(1, 10240, 2560).kind == "direct" - with pytest.raises(ValueError): - op(a.contiguous(), b) - a, b = _padded(8, 2560, 6144) - assert plan(8, 2560, 6144).kind == "splitk" - with pytest.raises(ValueError): - op(a, b.t().contiguous().t()) - - -def test_direct_branch_rejects_a_row_major_mat_b_loudly() -> None: - a = torch.randn(1, 2560, dtype=torch.bfloat16, device="cuda") - b = torch.randn(2560, 10240, dtype=torch.bfloat16, device="cuda") - assert plan(1, 10240, 2560).kind == "direct" - with pytest.raises(ValueError): - op(a, b) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_gated_residual.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_gated_residual.py index 4fa6444fc13d..9d2a95df0a97 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_gated_residual.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_gated_residual.py @@ -4,8 +4,9 @@ The thing under test is a formula -- which operand is gated, which axis is reduced, where each division by `hc_count` sits, the factor on the injection --- and every mis-reading of it keeps the shapes. So beyond the cells, one test -measures how far outside the entry's band each wrong variant lands. +-- and every mis-reading of it keeps the shapes; the cells' band sits 32x..192x +inside the nearest wrong formula (measured). Beyond the cells: graph replay and +one loud misuse. """ import zlib @@ -13,7 +14,6 @@ import pytest import torch -import torch.nn.functional as F from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501 gate_ratio, @@ -127,104 +127,11 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res _compare((mixed, inject), op.reference(*args)) -def test_stack_final_mixer_matches_the_injecting_call() -> None: - """The injection branch is independent of `mixed`, bitwise.""" - _, (normed, down, up, inject_w, c) = _build((256,), seed=1) - mixed_only, none_weights = op(normed, down, up, None, c) - assert none_weights is None - with_inject, _ = op(normed, down, up, inject_w, c) - assert torch.equal(mixed_only, with_inject) - - -def test_flattened_call_matches_batched() -> None: - _, (normed, down, up, inject_w, c) = _build((3, 17), seed=2) - batched_mixed, batched_inject = op(normed, down, up, inject_w, c) - flat_mixed, flat_inject = op(normed.reshape(-1, WIDTH), down, up, inject_w, c) - assert torch.equal(flat_mixed, batched_mixed.reshape(-1, HIDDEN)) - assert torch.equal(flat_inject, batched_inject.reshape(-1, HC_COUNT)) - - -def test_far_above_the_checkpoint_range_the_outputs_stay_bounded() -> None: - """Finite, `|mixed| <= max|normed|`, `inject` in [0, 2], to 4096x the - stream. The float64 comparison is not the question up there: the sigmoid - saturates and bf16 rounding of its logit flips an all-but-binary gate.""" - _, (base, down, up, inject_w, c) = _build((256,), seed=6) - assert base.float().abs().max().item() > 20.0, "operands are not at the checkpoint's range" - for scale in (1.0, 8.0, 64.0, 512.0, 4096.0): - # Scale the normed stream: an RMS norm would undo a scaled input. - normed = (base.float() * scale).to(torch.bfloat16) - mixed, inject = op(normed, down, up, inject_w, c) - assert torch.isfinite(mixed.float()).all(), scale - assert torch.isfinite(inject.float()).all(), scale - bound = normed.float().abs().max().item() * (1.0 + torch.finfo(torch.bfloat16).eps) - assert mixed.float().abs().max().item() <= bound, scale - assert (inject.float() >= 0.0).all() and (inject.float() <= 2.0).all(), scale - - -def test_wrong_variants_land_outside_the_band() -> None: - """Each mis-reading keeps the shapes; each must land > 5x outside the band - the correct result sits inside.""" - raw, args = _build((512,), seed=7) - normed, down, up, inject_w, c = args - want_mixed, want_inject = op.reference(*args) - got_mixed, got_inject = op(*args) - print(f"\n correct mixed {gate_ratio(got_mixed, want_mixed):.4g}x the band") - print(f" correct inject {gate_ratio(got_inject, want_inject):.4g}x the band") - - x = normed.double() - streams = x.unflatten(-1, (c, HIDDEN)) - logits = x @ down.double().t() - gate = torch.sigmoid(F.silu(logits / c) @ up.double().t()).unflatten(-1, (c, HIDDEN)) - inject_logits = x @ inject_w.double().t() - wrong_mixed = { - "sum instead of mean": lambda: (gate * streams).sum(dim=-2), - "no /hc_count before silu": lambda: ( - torch.sigmoid(F.silu(logits) @ up.double().t()).unflatten(-1, (c, HIDDEN)) * streams - ).mean(dim=-2), - "silu and sigmoid swapped": lambda: ( - F.silu(torch.sigmoid(logits / c) @ up.double().t()).unflatten(-1, (c, HIDDEN)) * streams - ).mean(dim=-2), - "stream axis taken last": lambda: ( - gate.transpose(-1, -2).reshape(-1, c, HIDDEN) * streams - ).mean(dim=-2), - "gates the raw stream": lambda: (gate * raw.double().unflatten(-1, (c, HIDDEN))).mean( - dim=-2 - ), - "mean over hidden": lambda: (gate * streams) - .mean(dim=-1) - .repeat_interleave(HIDDEN // c, dim=-1), - } - wrong_inject = { - "missing the factor of 2": lambda: torch.sigmoid(inject_logits / c), - "no /hc_count on the logit": lambda: 2.0 * torch.sigmoid(inject_logits), - "tanh instead of 2*sigmoid": lambda: torch.tanh(inject_logits / c), - "silu instead of 2*sigmoid": lambda: F.silu(inject_logits / c), - } - for want, variants in ((want_mixed, wrong_mixed), (want_inject, wrong_inject)): - for name, build in variants.items(): - ratio = gate_ratio(build().to(torch.bfloat16), want) - print(f" {name:<28} {ratio:10.4g}x the band") - assert ratio > 5.0, f"the band does not separate {name!r} ({ratio:.4g}x)" - with pytest.raises(AssertionError): - _compare((build().to(torch.bfloat16), None), (want, None)) - - def test_malformed_operands_raise_rather_than_return() -> None: - """Why there is no `is_valid`: each of these is loud on its own.""" + """Why there is no `is_valid`: a width the weights do not take is loud on its own.""" _, (normed, down, up, inject_w, c) = _build((32,), seed=9) with pytest.raises(RuntimeError): op(normed[:, : WIDTH - 8], down, up, inject_w, c) - with pytest.raises(RuntimeError): - op(normed, down.float(), up, inject_w, c) - odd = WIDTH + 1 - with pytest.raises(RuntimeError): - op( - torch.randn(32, odd, dtype=torch.bfloat16, device=DEV), - torch.randn(LOWRANK, odd, dtype=torch.bfloat16, device=DEV), - torch.randn(odd, LOWRANK, dtype=torch.bfloat16, device=DEV), - torch.randn(3, odd, dtype=torch.bfloat16, device=DEV), - 3, - ) def test_cuda_graph_capture_and_replay() -> None: diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py index 2ca906f8b1e1..2cd74f1c0b3d 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py @@ -6,8 +6,8 @@ combine's and the finalize's own rounding at the MoE sites), `part` within `2^-16` of the reference relative to its scale, `ss` to fp32 accuracy. Beyond the cells: a stale-read check under CUDA-graph replay with PDL, where a kernel that read its -activations before `griddepcontrol.wait` would see the previous replay's; determinism; -the guards. +activations before `griddepcontrol.wait` would see the previous replay's; the one-row +guard. """ import zlib @@ -156,12 +156,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by _compare_moe(args, out, op.reference(*args)) -def test_deterministic() -> None: - args = _build(dict(moe=True, fold=10), seed=3) - for a, b in zip(op(*args), op(*args)): - assert torch.equal(a, b) - - def test_reads_after_the_wait_under_graph_replay() -> None: """A producer kernel rewrites every activation inside the graph right before this op. With PDL the op may start before the producer finishes; reading any activation before @@ -200,6 +194,3 @@ def test_guards() -> None: two = args[0].expand(2, -1).contiguous() with validating(op), pytest.raises(AssertionError, match="one row"): op(two, args[1].expand(2, -1).contiguous(), args[2].expand(2, -1).contiguous(), *args[3:]) - args[4] = args[4].contiguous() - with validating(op), pytest.raises(AssertionError): - op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py index b92d780b5cde..5958e069eb6a 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py @@ -5,8 +5,8 @@ of the CuTe mix kernel. Each entry's cells against its float64 reference in its band (`mixed` and `inject_out` in `rtol = eps`, `atol = eps * max|ref|`; for the NVFP4 build, the NVFP4 exact on the returned `mixed` outside the reciprocal tie window), a -determinism check, a stale-read check under CUDA-graph replay with PDL for each build, -the guards.""" +stale-read check under CUDA-graph replay with PDL for the NVFP4 build, the one-row +guard.""" import zlib @@ -120,20 +120,12 @@ def test_certified_cells(name, cell, certifying) -> None: # noqa: F811 — fixt _check_quant(out[0], out[2], out[3], args[8]) -def test_deterministic() -> None: - op = qwen4_exp_hc_decode_mix_fp4_cute - args = _build(dict(), seed=3) - for a, b in zip(op(*args), op(*args)): - assert torch.equal(a, b) - - -@pytest.mark.parametrize("name", list(ENTRIES)) -def test_reads_after_the_wait_under_graph_replay(name) -> None: +def test_reads_after_the_wait_under_graph_replay() -> None: """Producer kernels rewrite `part`, `ss` and `hc_out` inside the graph right before the op; with PDL a read before `griddepcontrol.wait` would see the previous replay's. - Each build is its own compiled variant, so each is replayed.""" - op, fp4 = ENTRIES[name], name == "fp4" - args = _build(dict(), seed=8, fp4=fp4) + The NVFP4 build: the MoE sites' variant, with the most outputs.""" + op = qwen4_exp_hc_decode_mix_fp4_cute + args = _build(dict(), seed=8) targets = (args[0], args[1], args[3]) sources = [t.clone() for t in targets] @@ -152,7 +144,7 @@ def step(): with torch.cuda.graph(graph): graph_out = step() for trial in range(3): - fresh = _build(dict(), seed=100 + trial, fp4=fp4) + fresh = _build(dict(), seed=100 + trial) for src, new in zip(sources, (fresh[0], fresh[1], fresh[3])): src.copy_(new) graph.replay() @@ -162,12 +154,9 @@ def step(): def test_guards() -> None: - """Both entries share the shape checks; the NVFP4 build's `is_valid` runs them.""" + """Both entries share the shape checks; the NVFP4 build's `is_valid` refuses two rows.""" op = qwen4_exp_hc_decode_mix_fp4_cute args = _build(dict(), seed=6) two = args[3].expand(2, -1).contiguous() with validating(op), pytest.raises(AssertionError): op(args[0].expand(-1, 2, -1), args[1].expand(-1, 2), args[2], two, *args[4:]) - contiguous_up = args[2].contiguous() - with validating(op), pytest.raises(AssertionError): - op(args[0], args[1], contiguous_up, *args[3:]) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first.py index 2bceb75ebecb..4ec79d79cbe3 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first.py @@ -1,15 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """GPU test for the qwen4_exp_hc_decode_tc_first entry: cells in the band against -the float64 reference, wrong formulas outside, determinism, the guard, graph -capture, and the decode mixer built on it against the gated-residual float64 -reference.""" +the float64 reference, the row-count guard, graph capture.""" import zlib import pytest import torch -import torch.nn.functional as F from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_decode_tc_first import ( # noqa: E501 qwen4_exp_hc_decode_tc_first as op, @@ -42,8 +39,7 @@ def _compare(out, ref) -> None: if got is None: continue assert got.dtype == want.dtype and got.shape == want.shape and got.is_contiguous() - if got.numel(): - assert _ratio(got, want) <= 1.0, f"{_ratio(got, want):.4g}x the band" + assert _ratio(got, want) <= 1.0, f"{_ratio(got, want):.4g}x the band" def _build(spec, seed): @@ -72,39 +68,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by _compare((low, inject), op.reference(*args)) -def test_deterministic() -> None: - args = _build(dict(rows=128), seed=7) - a, b = op(*args), op(*args) - assert torch.equal(a[0], b[0]) and torch.equal(a[1], b[1]) - - -def test_wrong_variants_land_outside_the_band() -> None: - normed, first, rank, c, wi = _build(dict(rows=64), seed=3) - want_low, want_inject = op.reference(normed, first, rank, c, wi) - got = op(normed, first, rank, c, wi) - lg = normed.double() @ first.double() - for name, value, want in ( - ("silu without /hc_count", F.silu(lg[:, :rank]), want_low), - ("sigmoid instead of silu", torch.sigmoid(lg[:, :rank] / c), want_low), - ("columns shifted by one", F.silu(lg[:, 1 : rank + 1] / c), want_low), - ("inject missing the 2", torch.sigmoid(lg[:, rank : rank + c] / c), want_inject), - ( - "inject from the padding", - 2 * torch.sigmoid(lg[:, rank + c : rank + 2 * c] / c), - want_inject, - ), - ): - ratio = _ratio(value.to(want.dtype), want) - print(f"\n {name:<28} {ratio:10.4g}x the band (op: {_ratio(got[0], want_low):.3g}x)") - assert ratio > 5.0, name - - -def test_guard_refuses_a_contiguous_weight() -> None: - normed, first, rank, c, wi = _build(dict(rows=4), seed=4) - with validating(op), pytest.raises(AssertionError): - op(normed, first.contiguous(), rank, c, wi) - - def test_guard_refuses_more_than_one_tile_of_rows() -> None: normed, first, rank, c, wi = _build(dict(rows=129), seed=4) with validating(op), pytest.raises(AssertionError): @@ -127,58 +90,3 @@ def test_cuda_graph_capture_and_replay() -> None: torch.cuda.synchronize() low, inject = op(fresh, first, rank, c, wi) assert torch.equal(g_low, low) and torch.equal(g_inject, inject) - - -def test_mixer_matches_the_gated_residual_reference() -> None: - """The decode mixer as `DecodeTarget._hc_mixer` runs it from 2 rows -- this op, - then `qwen4_exp_hc_up_gate_mean` -- against `qwen4_exp_gated_residual.reference` - (float64 from `normed`) in the band that entry certifies its own op in, beside - the cuBLAS + `qwen4_exp_hc_lowrank_act` mixer it replaces. One representative - case (128 rows, with injection): each component certifies its own cells.""" - rows, with_inject = 128, True - from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm - from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501 - gate_ratio, - qwen4_exp_gated_residual, - ) - from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_lowrank_act import ( # noqa: E501 - qwen4_exp_hc_lowrank_act, - ) - from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_up_gate_mean import ( # noqa: E501 - qwen4_exp_hc_up_gate_mean, - ) - - hc, hidden, eps = 4, 2560, 1e-6 - width = hc * hidden - g = torch.Generator(device=DEV).manual_seed(rows) - - def normal(*shape, scale): - return (torch.randn(*shape, generator=g, device=DEV) * scale).to(torch.bfloat16) - - # This checkpoint's layer-0 hyper-connection element statistics. - down, up = normal(RANK, width, scale=0.0296), normal(width, RANK, scale=0.0681) - inject_w, gamma = normal(hc, width, scale=0.0107), normal(width, scale=0.477) - gamma[0], gamma[width // 2] = 6.625, -6.625 - x = torch.randn(rows, hc, hidden, generator=g, device=DEV) - x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) - normed = (x.flatten(-2) * (1.0 + gamma.float())).to(torch.bfloat16) - first = torch.cat([down, inject_w], 0) if with_inject else down - first = torch.cat( - [first, torch.zeros(384 - first.shape[0], width, device=DEV, dtype=first.dtype)] - ) - low, inject = op(normed, first.t(), RANK, hc, with_inject) - mixed = qwen4_exp_hc_up_gate_mean(low, up.t(), normed, hc) - want = qwen4_exp_gated_residual.reference( - normed, down, up, inject_w if with_inject else None, hc - ) - assert gate_ratio(mixed, want[0]) <= 1.0, f"mixed {gate_ratio(mixed, want[0]):.3f}x the band" - if with_inject: - assert gate_ratio(inject, want[1]) <= 1.0, f"inject {gate_ratio(inject, want[1]):.3f}x" - old_low, _ = qwen4_exp_hc_lowrank_act( - cublas_mm(normed, first.t(), None, torch.float32), RANK, hc, with_inject, torch.bfloat16 - ) - prev = qwen4_exp_hc_up_gate_mean(old_low, up.t(), normed, hc) - print( - f"\n rows {rows}: tc first {gate_ratio(mixed, want[0]):.3f}x band, " - f"cublas + lowrank_act {gate_ratio(prev, want[0]):.3f}x" - ) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_inject_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_inject_norm.py index 0341002b3ddf..b69e82f5c120 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_inject_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_inject_norm.py @@ -4,9 +4,8 @@ Cells are gated at one ulp of the activation dtype against the float64 reference; the MoE-combine cells widen the residual's band by the combine's own -rounding (see `_compare_moe`). Beyond them: agreement with the compositions it -replaces, the separation from plausible wrong formulas, the norm-only aliasing -claim, and graph capture. +rounding (see `_compare_moe`). Beyond them: the CuTe fold kernel bitwise the +Triton kernel at the last row count before the switch, and the slot-map guard. """ import sys @@ -18,18 +17,6 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_inject_norm import ( # noqa: E501 qwen4_exp_hc_inject_norm as op, ) -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, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp4_block_scale_moe_renorm_permuted import ( # noqa: E501 - finalize, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fused_sigmoid_gate_mul_add import ( # noqa: E501 - fused_sigmoid_gate_mul_add, -) -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, -) __extra_import_path__ = ["../../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name @@ -186,208 +173,19 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by _compare_moe(args, (hc, normed), op.reference(*args)) -def _old_composition(hyper_input, block_out, inject, gamma, eps, hc_count): - hc = qwen4_exp_post_inject(hyper_input, block_out, inject) - hidden = hc.shape[1] // hc_count - parts = torch.split(hc, [hidden] * hc_count, 1) - weights = torch.split(gamma, hidden, 0) - normed = torch.cat([flashinfer_gemma_rmsnorm(p, w, eps) for w, p in zip(weights, parts)], 1) - return hc, normed - - -def test_matches_the_composition_it_replaces() -> None: - """Against post_inject + four gemma norms + cat: the residual within one - ulp (fma contraction only), the normed stream within one ulp. One - representative prefill chunk: each component certifies its own cells.""" - rows = 8192 - args = _build(dict(rows=rows), seed=rows)[:6] - hc, normed = op(*args) - old_hc, old_normed = _old_composition(*args) - assert _ulp_ratio(hc, old_hc) <= 1.0 - assert _ulp_ratio(normed, old_normed) <= 1.0 - print(f"\n rows {rows}: hc bitwise-equal fraction {(hc == old_hc).float().mean().item():.6f}") - - -def test_wrong_variants_land_outside_the_band() -> None: - """Each mis-reading keeps the shapes; each lands > 5x outside one ulp. The - streams carry different scales, as the trained residual's do, so a norm - over the wrong group is visible.""" - hyper_input, block_out, inject, gamma, eps, c, _, _ = _build(dict(rows=64), seed=5) - scales = torch.tensor([0.5, 1.0, 2.0, 4.0], device=DEV).repeat_interleave(HIDDEN) - hyper_input = (hyper_input.float() * scales).to(torch.bfloat16) - want_hc, want_normed = op.reference(hyper_input, block_out, inject, gamma, eps, c) - x = hyper_input.double().unflatten(-1, (c, HIDDEN)) - b, g = block_out.double(), inject.double() - wrong_hc = { - "stream-minor inject": (x.flatten(-2) + (b.unsqueeze(-1) * g.unsqueeze(-2)).flatten(-2)), - "no gate": (x + b.unsqueeze(-2)).flatten(-2), - } - for name, value in wrong_hc.items(): - ratio = _ulp_ratio(value.to(torch.bfloat16), want_hc) - print(f"\n hc {name:<24} {ratio:10.4g} ulp") - assert ratio > 5.0, name - hc = want_hc.double().unflatten(-1, (c, HIDDEN)) - full = hc.flatten(-2) - wrong_normed = { - "RMS over the full width": full - * torch.rsqrt(full.pow(2).mean(-1, keepdim=True) + eps) - * (1.0 + gamma.double()), - "w instead of 1 + w": (hc * torch.rsqrt(hc.pow(2).mean(-1, keepdim=True) + eps)).flatten(-2) - * gamma.double(), - "norm of the pre-inject stream": (x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)) - .flatten(-2) - .mul(1.0 + gamma.double()), - } - for name, value in wrong_normed.items(): - ratio = _ulp_ratio(value.to(torch.bfloat16), want_normed) - print(f" normed {name:<30} {ratio:10.4g} ulp") - assert ratio > 5.0, name - - -def test_guard_refuses_block_out_without_inject() -> None: - hyper_input, block_out, inject, gamma, eps, c, _, _ = _build(dict(rows=4), seed=6) - with validating(op), pytest.raises(AssertionError): - op(hyper_input, block_out, None, gamma, eps, c) - with validating(op), pytest.raises(AssertionError): - op(hyper_input, block_out, inject[0], gamma, eps, c) - - -def test_cuda_graph_capture_and_replay() -> None: - args = list(_build(dict(rows=8), seed=8)) - static_in = args[0] - 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): - graph_hc, graph_normed = op(*args) - fresh = _build(dict(rows=8), seed=18)[0] - static_in.copy_(fresh) - graph.replay() - torch.cuda.synchronize() - eager_hc, eager_normed = op(fresh, *args[1:]) - assert torch.equal(graph_hc, eager_hc) - assert torch.equal(graph_normed, eager_normed) - - -@pytest.mark.parametrize("rows", [128, 2048]) -def test_fold_matches_the_runner_finalize(rows) -> None: - """With the fold, bitwise this op on the runner's finalized rows: `finalize` - (the permuted MoE entry's emulation of the runner's finalize, bitwise its - output) and then the combine without the fold. 128 rows run the CuTe kernel, - 2048 the Triton kernel's two-warp form.""" - args = _build(dict(rows=rows, moe=True, fold=10), seed=300 + rows) - hyper_input, rows_, inject, gamma, eps, c, shared, gate, topk_weights, topk_rows = args - hc, normed = op(*args) - routed = finalize(rows_, topk_weights, topk_rows) - old_hc, old_normed = op(hyper_input, routed, inject, gamma, eps, c, shared, gate) - assert torch.equal(hc, old_hc) - assert torch.equal(normed, old_normed) - - -def test_moe_combine_matches_the_composition_it_replaces() -> None: - """Against `fused_sigmoid_gate_mul_add` (in place on the routed output, as - the forward called it) and then this op without the combine: bitwise, at one - representative decode batch.""" - rows = 128 - args = _build(dict(rows=rows, moe=True), seed=100 + rows) - hyper_input, block_out, inject, gamma, eps, c, shared, gate = args - hc, normed = op(*args) - combined = fused_sigmoid_gate_mul_add(block_out.clone(), gate, shared) - old_hc, old_normed = op(hyper_input, combined, inject, gamma, eps, c) - print( - f"\n rows {rows}: hc bitwise {(hc == old_hc).float().mean().item():.6f}, " - f"normed bitwise {(normed == old_normed).float().mean().item():.6f}" - ) - assert torch.equal(hc, old_hc) - assert torch.equal(normed, old_normed) - - -def test_moe_wrong_variants_land_outside_the_band() -> None: - """Each mis-reading of the combine keeps the shapes; each lands > 5x outside - the `_compare_moe` residual band.""" - args = _build(dict(rows=64, moe=True), seed=7) - hyper_input, block_out, inject, gamma, eps, c, shared, gate = args - want_hc = op.reference(*args)[0] - feps = torch.finfo(torch.bfloat16).eps - b = _combine(block_out, shared, gate).double().abs() - band = feps * want_hc.double().abs() + ( - feps * b.unsqueeze(-2) * inject.double().abs().unsqueeze(-1) - ).flatten(-2) - r, sh, s = block_out.double(), shared.double(), torch.sigmoid(gate.double()) - variants = { - "no sigmoid": r + gate.double() * sh, - "gate on the routed half": s * r + sh, - "gate dropped": r + sh, - "shared dropped": r, - "1 - sigmoid": r + (1.0 - s) * sh, - } - x = hyper_input.double().unflatten(-1, (c, HIDDEN)) - for name, blk in variants.items(): - wrong = (x + blk.unsqueeze(-2) * inject.double().unsqueeze(-1)).flatten(-2) - ratio = ((wrong.to(torch.bfloat16).double() - want_hc.double()).abs() / band).max().item() - print(f"\n hc {name:<26} {ratio:10.4g}x the band") - assert ratio > 5.0, name - - -def test_guard_refuses_a_partial_combine() -> None: - hyper_input, block_out, inject, gamma, eps, c, shared, gate = _build( - dict(rows=4, moe=True), seed=9 - ) +def test_guard_refuses_an_int64_slot_map() -> None: + """The fold's slot map is the runner's int32 index map; `is_valid` refuses an int64 one.""" fold = _build(dict(rows=4, moe=True, fold=10), seed=9) - with validating(op), pytest.raises(AssertionError): - op(*fold[:6], None, None, *fold[8:]) - with validating(op), pytest.raises(AssertionError): - op(*fold[:9], None) - with validating(op), pytest.raises(AssertionError): + with validating(op), pytest.raises(AssertionError, match="int32"): op(*fold[:9], fold[9].long()) - with validating(op), pytest.raises(AssertionError): - op(hyper_input, block_out, inject, gamma, eps, c, shared, None) - with validating(op), pytest.raises(AssertionError): - op(hyper_input, None, None, gamma, eps, c, shared, gate) - with validating(op), pytest.raises(AssertionError): - op(hyper_input, block_out, inject, gamma, eps, c, shared, gate[:1].expand(4, 1)) - -def test_moe_cuda_graph_capture_and_replay() -> None: - args = list(_build(dict(rows=8, moe=True), seed=10)) - 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): - graph_hc, graph_normed = op(*args) - fresh = _build(dict(rows=8, moe=True), seed=20) - for i in (0, 1, 6, 7): - args[i].copy_(fresh[i]) - graph.replay() - torch.cuda.synchronize() - eager_hc, eager_normed = op(*args) - assert torch.equal(graph_hc, eager_hc) - assert torch.equal(graph_normed, eager_normed) - -@pytest.mark.parametrize( - "spec", - [ - dict(rows=2, drop=True), - dict(rows=128), - dict(rows=511), - ], - ids=lambda s: "-".join(f"{k}{v}" for k, v in s.items()), -) -def test_fold_cute_is_the_triton_kernel_bitwise(spec, monkeypatch) -> None: +def test_fold_cute_is_the_triton_kernel_bitwise(monkeypatch) -> None: """Below `_FEW_WARPS_MIN_ROWS` rows the fold runs the CuTe DSL kernel; its `hc_out` and `normed` are bitwise the Triton kernel's (run by refusing the - CuTe guard): the smallest decode batch with dropped slots, the largest decode - batch, and the last row count before the switch.""" - args = _build(dict(spec, moe=True, fold=10), seed=500 + spec["rows"]) + CuTe guard) at the last row count before the switch.""" + rows = M._FEW_WARPS_MIN_ROWS - 1 + args = _build(dict(rows=rows, moe=True, fold=10), seed=500 + rows) calls = [] cute_fold = M._fold_cute @@ -397,16 +195,8 @@ def recording(*a): monkeypatch.setattr(M, "_fold_cute", recording) hc, normed = op(*args) - assert calls == [spec["rows"]], "the decode MoE site runs the CuTe kernel" + assert calls == [rows], "the decode MoE site runs the CuTe kernel" monkeypatch.setattr(M, "_fold_cute_ok", lambda *a: False) want_hc, want_normed = op(*args) assert torch.equal(hc, want_hc) assert torch.equal(normed, want_normed) - - -def test_fold_from_the_few_warps_rows_keeps_triton(monkeypatch) -> None: - """From `_FEW_WARPS_MIN_ROWS` rows (the Triton kernel's two-warp form, another - reduction order) the fold stays on the Triton kernel.""" - args = _build(dict(rows=M._FEW_WARPS_MIN_ROWS, moe=True, fold=10), seed=77) - monkeypatch.setattr(M, "_fold_cute", lambda *a: pytest.fail("CuTe kernel at 512 rows")) - op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_lowrank_act.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_lowrank_act.py index 8226d441001d..8641b2b8b7b5 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_lowrank_act.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_lowrank_act.py @@ -1,13 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """GPU test for the qwen4_exp_hc_lowrank_act entry: cells within one ulp of the -float64 reference, wrong formulas outside, graph capture.""" +float64 reference, the fp32-logits guard, graph capture.""" import zlib import pytest import torch -import torch.nn.functional as F from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_lowrank_act import ( # noqa: E501 qwen4_exp_hc_lowrank_act as op, @@ -37,8 +36,7 @@ def _compare(out, ref) -> None: if got is None: continue assert got.dtype == want.dtype and got.shape == want.shape and got.is_contiguous() - if got.numel(): - assert _ulp_ratio(got, want) <= 1.0, f"{_ulp_ratio(got, want):.4g} ulp" + assert _ulp_ratio(got, want) <= 1.0, f"{_ulp_ratio(got, want):.4g} ulp" def _build(spec, seed): @@ -58,25 +56,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by _compare((low, inject), op.reference(*args)) -def test_wrong_variants_land_outside_the_band() -> None: - logits, rank, c, wi, dt = _build(dict(rows=64), seed=3) - want_low, want_inject = op.reference(logits, rank, c, wi, dt) - lg = logits.double() - for name, value, want in ( - ("silu without /hc_count", F.silu(lg[:, :rank]), want_low), - ("sigmoid instead of silu", torch.sigmoid(lg[:, :rank] / c), want_low), - ("inject missing the 2", torch.sigmoid(lg[:, rank : rank + c] / c), want_inject), - ( - "inject from the padding", - 2 * torch.sigmoid(lg[:, rank + c : rank + 2 * c] / c), - want_inject, - ), - ): - ratio = _ulp_ratio(value.to(dt), want) - print(f"\n {name:<28} {ratio:10.4g} ulp") - assert ratio > 5.0, name - - def test_guard_refuses_bf16_logits() -> None: logits, rank, c, wi, dt = _build(dict(rows=4), seed=4) with validating(op), pytest.raises(AssertionError): diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py index cea97275bf3a..3e392c8e6f75 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py @@ -1,9 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """GPU test for the qwen4_exp_hc_up_gate_mean entry: cells in the band -`qwen4_exp_gated_residual`'s test derives for `mixed`, wrong formulas outside, -the guard, graph capture, and the mixer as the forward runs it against the -gated-residual float64 reference.""" +`qwen4_exp_gated_residual`'s test derives for `mixed`, the guard, graph +capture.""" import zlib @@ -29,8 +28,7 @@ def _compare(out, ref) -> None: """`rtol = eps`, `atol = eps * max|ref|` (`gate_ratio <= 1`), the band `qwen4_exp_gated_residual` derives for `mixed`.""" assert out.dtype == ref.dtype and out.shape == ref.shape and out.is_contiguous() - if out.numel(): - assert gate_ratio(out, ref) <= 1.0, f"{gate_ratio(out, ref):.4g}x the band" + assert gate_ratio(out, ref) <= 1.0, f"{gate_ratio(out, ref):.4g}x the band" def _build(spec, seed): @@ -56,25 +54,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by _compare(mixed, op.reference(*args)) -def test_wrong_variants_land_outside_the_band() -> None: - low, up, normed, c = _build(dict(rows=128), seed=3) - want = op.reference(low, up, normed, c) - got = op(low, up, normed, c) - h = normed.shape[1] // c - logits = (low.double() @ up.double()).unflatten(-1, (c, h)) - x = normed.double().unflatten(-1, (c, h)) - g = torch.sigmoid(logits) - for name, value in ( - ("sum instead of mean", (g * x).sum(-2)), - ("no sigmoid", (logits * x).mean(-2)), - ("stream-minor packing", (g.transpose(-1, -2).reshape(-1, c, h) * x).mean(-2)), - ("ungated mean", x.mean(-2)), - ): - ratio = gate_ratio(value.to(want.dtype), want) - print(f"\n {name:<24} {ratio:10.4g}x the band (op: {gate_ratio(got, want):.3g}x)") - assert ratio > 5.0, name - - def test_guard_refuses_a_contiguous_up() -> None: low, up, normed, c = _build(dict(rows=4), seed=4) with validating(op), pytest.raises(AssertionError): @@ -97,45 +76,3 @@ def test_cuda_graph_capture_and_replay() -> None: graph.replay() torch.cuda.synchronize() assert torch.equal(g_mixed, op(fresh[0], up, fresh[2], c)) - - -def test_mixer_matches_the_gated_residual_reference() -> None: - """The mixer as the forward runs it -- `cublas_mm` over the zero-padded - `[down; inject]` (fp32 out), `qwen4_exp_hc_lowrank_act`, this op -- against - `qwen4_exp_gated_residual.reference` (float64 from `normed`), in the band - that entry certifies its own op in. One prefill-sized case: each component - is certified over its own cells, and the injection columns never reach - `mixed`.""" - from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm - from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501 - qwen4_exp_gated_residual, - ) - from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_lowrank_act import ( # noqa: E501 - qwen4_exp_hc_lowrank_act, - ) - - rows, hc, hidden, rank, eps = 2048, 4, 2560, 320, 1e-6 - width = hc * hidden - g = torch.Generator(device=DEV).manual_seed(rows) - - def normal(*shape, scale): - return (torch.randn(*shape, generator=g, device=DEV) * scale).to(torch.bfloat16) - - # This checkpoint's layer-0 hyper-connection element statistics. - down, up = normal(rank, width, scale=0.0296), normal(width, rank, scale=0.0681) - inject_w, gamma = normal(hc, width, scale=0.0107), normal(width, scale=0.477) - gamma[0], gamma[width // 2] = 6.625, -6.625 - x = torch.randn(rows, hc, hidden, generator=g, device=DEV) - x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) - normed = (x.flatten(-2) * (1.0 + gamma.float())).to(torch.bfloat16) - first = torch.cat([down, inject_w], 0) - first = torch.cat( - [first, torch.zeros(384 - first.shape[0], width, device=DEV, dtype=first.dtype)] - ) - low, _ = qwen4_exp_hc_lowrank_act( - cublas_mm(normed, first.t(), None, torch.float32), rank, hc, True, torch.bfloat16 - ) - mixed = op(low, up.t(), normed, hc) - want = qwen4_exp_gated_residual.reference(normed, down, up, inject_w, hc) - _compare(mixed, want[0]) - print(f"\n rows {rows}: fused {gate_ratio(mixed, want[0]):.3f}x band") diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4.py index 17cb0509f5ff..eecc9d836bef 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4.py @@ -1,17 +1,15 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """GPU test for the qwen4_exp_hc_up_gate_mean_fp4 entry: `mixed` in the band -`qwen4_exp_hc_up_gate_mean`'s test uses and bitwise that op's, the NVFP4 -outputs bitwise against `torch.ops.trtllm.fp4_quantize` of the returned -`mixed` and against the native-torch quantization outside the reciprocal tie -window, wrong recipes outside, guards, graph capture.""" +`qwen4_exp_hc_up_gate_mean`'s test uses, the NVFP4 outputs against the +native-torch quantization of the returned `mixed` outside the reciprocal tie +window, the global-scale guard, graph capture.""" import zlib import pytest import torch -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( Fp4Quantize, fp4_quantize, @@ -19,9 +17,6 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501 gate_ratio, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_up_gate_mean import ( # noqa: E501 - qwen4_exp_hc_up_gate_mean, -) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_up_gate_mean_fp4 import ( # noqa: E501 qwen4_exp_hc_up_gate_mean_fp4 as op, ) @@ -69,9 +64,6 @@ def _compare(out, ref, gs) -> None: (mixed, xq, xsf), (ref_mixed, _, _) = out, ref assert mixed.dtype == ref_mixed.dtype and mixed.shape == ref_mixed.shape assert mixed.is_contiguous() - if mixed.numel() == 0: - assert xq.numel() == 0 and xsf.numel() == 0 - return assert gate_ratio(mixed, ref_mixed) <= 1.0, f"{gate_ratio(mixed, ref_mixed):.4g}x the band" _check_quant(mixed, xq, xsf, gs) @@ -105,73 +97,11 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by _compare(out, op.reference(*args), args[4]) -@pytest.mark.parametrize("rows", [2, 16390]) -def test_bitwise_the_composition_it_replaces(rows) -> None: - """`qwen4_exp_hc_up_gate_mean` then `torch.ops.trtllm.fp4_quantize` (linear - scales): every output bitwise equal, ties included. A one-tile decode batch - and a many-tile one with a ragged tail; the cells certify the rest.""" - low, up, normed, c, gs = _build(dict(rows=rows), seed=rows) - mixed, xq, xsf = op(low, up, normed, c, gs) - old = qwen4_exp_hc_up_gate_mean(low, up, normed, c) - assert torch.equal(mixed, old) - old_q, old_sf = torch.ops.trtllm.fp4_quantize(old, gs, 16, False, False) - assert torch.equal(xq, old_q.view(torch.uint8)), ( - f"{(xq != old_q.view(torch.uint8)).sum().item()} bytes" - ) - assert torch.equal(xsf, old_sf.view(torch.uint8).reshape(-1)) - - -@pytest.mark.parametrize("scale", [1e-6, 448.0 * 6.0 / 24.0, 1e4]) -def test_bitwise_across_global_scales(scale) -> None: - """A tiny scale (zero / subnormal e4m3 bytes), the calibrated one and a - saturating one.""" - low, up, normed, c, _ = _build(dict(rows=128), seed=11) - h = normed.shape[1] // c - for s in range(c): # row 0's first 64 mix columns are exactly zero: amax 0 - normed[0, s * h : s * h + 64] = 0 - gs = torch.tensor([scale], device=DEV, dtype=torch.float32) - mixed, xq, xsf = op(low, up, normed, c, gs) - assert not mixed[0, :64].any() - old_q, old_sf = torch.ops.trtllm.fp4_quantize(mixed, gs, 16, False, False) - assert torch.equal(xq, old_q.view(torch.uint8)) - assert torch.equal(xsf, old_sf.view(torch.uint8).reshape(-1)) - - -def test_wrong_recipes_fail_the_quant_check() -> None: - """Each plausible wrong layout or recipe is refused by `_check_quant`.""" - low, up, normed, c, gs = _build(dict(rows=128), seed=3) - mixed, xq, xsf = op(low, up, normed, c, gs) - _check_quant(mixed, xq, xsf, gs) - rows = mixed.shape[0] - swapped = ((xq & 0xF) << 4) | (xq >> 4) - sf2 = xsf.view(rows, -1) - # 128 rows: the 128x4-swizzled buffer has exactly the linear one's length. - sw = torch.ops.trtllm.fp4_quantize(mixed, gs, 16, False, True)[1].view(torch.uint8).reshape(-1) - h = normed.shape[1] // c - unrounded = ( - torch.sigmoid(low.float() @ up.float()).unflatten(-1, (c, h)) - * normed.float().unflatten(-1, (c, h)) - ).mean(-2) - q32, sf32 = nvfp4_reference(unrounded, gs) - for name, args in ( - ("nibbles swapped", (mixed, swapped, xsf)), - ("swizzled scales", (mixed, xq, sw)), - ("scale one code high", (mixed, xq, (sf2 + 1).reshape(-1))), - ("quantized before the bf16 rounding", (mixed, q32, sf32)), - ): - with pytest.raises(AssertionError): - _check_quant(*args, gs) - print(f"\n {name:<36} refused") - - def test_guards() -> None: + """The kernel loads `global_scale[0]` only; a per-row scale is refused.""" low, up, normed, c, gs = _build(dict(rows=4), seed=4) - with validating(op), pytest.raises(AssertionError): - op(low, up.contiguous(), normed, c, gs) with validating(op), pytest.raises(AssertionError): op(low, up, normed, c, torch.cat([gs, gs])) - with validating(op), pytest.raises(AssertionError): - op(low, up, normed, c, gs.double()) def test_cuda_graph_capture_and_replay() -> None: diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_post_inject.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_post_inject.py index 65853efa4a69..e609e8632440 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_post_inject.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_post_inject.py @@ -2,9 +2,8 @@ # SPDX-License-Identifier: Apache-2.0 """GPU test for the qwen4_exp_post_inject entry. -The cells are gated bitwise (16-bit) by `_compare`. Beyond them: -the separation from the plausible wrong operators, the bitwise identity and -one-hot claims, the leading-dim guard, and graph capture. +The cells are gated bitwise (16-bit) by `_compare`. Beyond them: the +leading-dim guard and graph capture. """ import zlib @@ -23,24 +22,18 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: - """Bitwise for 16-bit operands: the product of two bf16 (or fp16) - values is exact in fp32, so op and reference round the same fp32 sum - once -- measured bit-exact at every token count. fp32 operands round the - product, and whether it is fused into the add is the kernel's choice: - torch's default band there.""" + """Bitwise for 16-bit operands: the product of two bf16 values is exact in + fp32, so op and reference round the same fp32 sum once -- measured + bit-exact at every token count.""" assert out.dtype == ref.dtype, (out.dtype, ref.dtype) assert out.shape == ref.shape, (out.shape, ref.shape) assert out.is_contiguous() - if ref.dtype in (torch.bfloat16, torch.float16): - assert torch.equal(out, ref), (out.float() - ref.float()).abs().max().item() - else: - torch.testing.assert_close(out, ref) + assert torch.equal(out, ref), (out.float() - ref.float()).abs().max().item() DEV = "cuda" HC_COUNT = 4 HIDDEN = 2560 -WIDTH = HC_COUNT * HIDDEN def _operands(leading, hc_count=HC_COUNT, hidden=HIDDEN, dtype=torch.bfloat16, seed=0): @@ -60,11 +53,7 @@ def _build(spec, seed): hyper_input, block_out, inject = _operands(spec["leading"], hc_count, hidden, dtype, seed) layout = spec.get("layout", "contiguous") width = hc_count * hidden - if layout == "row_slice": - buf = torch.randn(2 * hyper_input.shape[0], width, device=DEV).to(dtype) - buf[::2] = hyper_input - hyper_input = buf[::2] - elif layout == "col_slice": + if layout == "col_slice": buf = torch.randn(hyper_input.shape[0], 2 * width, device=DEV).to(dtype) buf[:, :width] = hyper_input hyper_input = buf[:, :width] @@ -85,104 +74,16 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res _compare(out, op.reference(hyper_input, block_out, inject)) -def test_rounds_once_in_fp32_not_twice_in_bf16() -> None: - """The bitwise gate is not the bf16-multiply-then-bf16-add form (0.0625 apart).""" - for num_tokens in (64, 2048): - hyper_input, block_out, inject = _operands((num_tokens,), seed=100 + num_tokens) - out = op(hyper_input, block_out, inject) - double_round = ( - hyper_input.unflatten(-1, (HC_COUNT, HIDDEN)) - + block_out.unsqueeze(-2) * inject.unsqueeze(-1) - ).flatten(-2) - assert torch.equal(out, op.reference(hyper_input, block_out, inject)) - assert (out.float() - double_round.float()).abs().max().item() > 0.0, ( - "the two rounding orders did not separate" - ) - - -def test_wrong_operators_are_separated() -> None: - """Stream-minor packing, no gate, gate on the residual: each > 1.0 away.""" - hyper_input, block_out, inject = _operands((64,), seed=3) - out = op(hyper_input, block_out, inject).float() - hi32, bo32, in32 = hyper_input.float(), block_out.float(), inject.float() - wrong = { - "stream_minor": ( - hi32.unflatten(-1, (HIDDEN, HC_COUNT)) + bo32.unsqueeze(-1) * in32.unsqueeze(-2) - ).flatten(-2), - "no_gate": (hi32.unflatten(-1, (HC_COUNT, HIDDEN)) + bo32.unsqueeze(-2)).flatten(-2), - "gate_on_residual": ( - hi32.unflatten(-1, (HC_COUNT, HIDDEN)) * in32.unsqueeze(-1) + bo32.unsqueeze(-2) - ).flatten(-2), - } - for name, variant in wrong.items(): - dev = (out - variant).abs().max().item() - print(f" {name:<18} max deviation {dev:.6g}") - assert dev > 1.0, (name, dev) - - -def test_zero_gate_is_the_identity() -> None: - hyper_input, block_out, _ = _operands((33,), seed=4) - inject = torch.zeros(33, HC_COUNT, dtype=torch.bfloat16, device=DEV) - assert torch.equal(op(hyper_input, block_out, inject), hyper_input) - - -def test_one_hot_gate_touches_exactly_one_stream() -> None: - """The property a token-broadcast bug destroys; asserted bitwise.""" - hyper_input, block_out, _ = _operands((16,), seed=5) - streams = hyper_input.unflatten(-1, (HC_COUNT, HIDDEN)) - for live in range(HC_COUNT): - inject = torch.zeros(16, HC_COUNT, dtype=torch.bfloat16, device=DEV) - inject[:, live] = 1.0 - out = op(hyper_input, block_out, inject).unflatten(-1, (HC_COUNT, HIDDEN)) - for s in range(HC_COUNT): - if s == live: - expected = (streams[:, s].float() + block_out.float()).to(torch.bfloat16) - assert torch.equal(out[:, s], expected), (live, s) - else: - assert torch.equal(out[:, s], streams[:, s]), (live, s) - - def test_guard_refuses_unequal_leading_dims() -> None: - """Unguarded, a 1-D `inject` or `block_out` broadcasts one row over every token.""" + """Unguarded, a 1-D `inject` broadcasts one row over every token.""" hyper_input, block_out, inject = _operands((8,), seed=10) unguarded = op(hyper_input, block_out, inject[0]) assert unguarded.shape == hyper_input.shape, "the hazard this guard covers is gone" assert not torch.equal(unguarded, op.reference(hyper_input, block_out, inject)) - unguarded = op(hyper_input, block_out[0], inject) - assert unguarded.shape == hyper_input.shape, "the hazard this guard covers is gone" - for args in ( - (hyper_input, block_out, inject[0]), - (hyper_input, block_out[0], inject), - (hyper_input, block_out[:4], inject), - (hyper_input, block_out, inject[:4]), - ): - with pytest.raises(AssertionError, match="leading dims must match"): - with validating(op): - op(*args) - - -def test_loud_refusals_are_not_guarded() -> None: - """A width that does not factor, and a host operand, raise on their own.""" - _, block_out, inject = _operands((8,), seed=11) - wrong = torch.randn(8, WIDTH + 16, dtype=torch.bfloat16, device=DEV) - with pytest.raises(RuntimeError, match="don't multiply up to the size of dim"): - op(wrong, block_out, inject) - hyper_input, _, inject = _operands((8,), seed=12) - with pytest.raises(RuntimeError, match="Expected all tensors to be on the same device"): - op(hyper_input, torch.randn(8, HIDDEN, dtype=torch.bfloat16), inject) - - -def test_mixed_dtypes_promote_the_output() -> None: - """Recorded, not guarded: an fp32 operand silently widens the residual stream.""" - hyper_input, block_out, _ = _operands((8,), seed=13) - fp32_inject = torch.rand(8, HC_COUNT, dtype=torch.float32, device=DEV) * 2.0 - out = op(hyper_input, block_out, fp32_inject) - assert out.dtype == torch.float32 - _compare(out, op.reference(hyper_input, block_out, fp32_inject)) - fp32_block_out = torch.randn(8, HIDDEN, dtype=torch.float32, device=DEV) - _, _, inject = _operands((8,), seed=13) - assert op(hyper_input, fp32_block_out, inject).dtype == torch.float32 + with pytest.raises(AssertionError, match="leading dims must match"): + with validating(op): + op(hyper_input, block_out, inject[0]) def test_cuda_graph_capture_and_replay() -> None: diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py index ea92991e4625..f23d2789ee81 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py @@ -5,10 +5,7 @@ The entry owns the contract. This file builds the checkpoint-shaped operands -- 512 experts in the CuteDSL weight layout, the activation as the target quantizes it, the sort as moe_sort returns it -- and holds the claims no cell carries: -per-pair invariance across tiles, tactics, row order and prefixes; CUDA-graph -replay at the captured row counts; the regions the op leaves unwritten; that -the gate rejects a wrong layout of each operand; the guards and the loud -rejections. +per-pair invariance across tiles and tactics, and the contiguity guard. """ import functools @@ -47,10 +44,8 @@ ACT_AMAX = 48.0 GSF = { "calibrated": 448 * 6 / ACT_AMAX, - "saturating": 25 * 448 * 6 / ACT_AMAX, "subnormal": 2.0**-12 * 448 * 6 / ACT_AMAX, } -SENTINEL = 0xA5 #: Tie windows of `_compare`, derived in its docstring. SF_TIE = 2.0**-20 @@ -59,13 +54,6 @@ #: (4 of them when that is more). FLIP_BUDGET = 1e-5 -#: Row counts of a decode capture (requests x MTP verify width 1..4, up to 32 -#: requests): both ends of the sort's block kernel, its first cluster-kernel T and -#: the largest capture. The eager prefix lengths straddle the captured ceiling and -#: reach the 8192-row one. -_CAPTURED_ROWS = (1, 4, 5, 128) -_PREFIXES = (1, 4, 5, 127, 128, 129, 4096, 8191) - # ── the gate ────────────────────────────────────────────────────────────── @@ -203,82 +191,61 @@ def _experts(num_experts: int = E): return w, s, ws2 -def _row_perm(n: int, layout: str) -> torch.Tensor: - """Destination row -> source row of the `[up; gate]` stack.""" +def _row_perm(n: int) -> torch.Tensor: + """Destination row -> source row of the `[up; gate]` stack, in `GATE_UP_GROUP`-row + groups.""" idx = torch.arange(n, device=DEV) - if layout == "group64": - return idx.view(2, n // (2 * GATE_UP_GROUP), GATE_UP_GROUP).transpose(0, 1).reshape(-1) - if layout == "rowwise": # trtllm-gen's interleave, without its 32-row shuffle - return idx.view(2, n // 2).transpose(0, 1).reshape(-1) - raise ValueError(layout) + return idx.view(2, n // (2 * GATE_UP_GROUP), GATE_UP_GROUP).transpose(0, 1).reshape(-1) @functools.lru_cache(maxsize=None) -def _weights(layout: str = "group64", num_experts: int = E): +def _weights(): """`(weight, weight_scale, alpha)` in the layout CuteDSL reads -- what the - target's load-time transform produces -- or one of the wrong layouts the - discrimination test feeds it.""" - w, s, ws2 = _experts(num_experts) - perm = _row_perm(2 * INTER, "group64" if layout == "scales_unpermuted" else layout) + target's load-time transform produces.""" + w, s, ws2 = _experts() + perm = _row_perm(2 * INTER) w_k = w[:, perm].contiguous() - s_rows = s if layout == "scales_unpermuted" else s[:, perm] - s_k = torch.ops.trtllm.block_scale_interleave(s_rows.contiguous()).view(s.shape) + s_k = torch.ops.trtllm.block_scale_interleave(s[:, perm].contiguous()).view(s.shape) alpha = (ISC1 * ws2).contiguous() return w_k.view(torch.float4_e2m1fn_x2), s_k, alpha -def _activation(tokens: int, seed: int, zero_tokens: int = 0, swizzled: bool = False): +def _activation(tokens: int, seed: int, zero_tokens: int = 0): """The target's FC1 input: bf16 hidden rows through `fp4_quantize(x, 1/isc1, 16, False, False)`, scales viewed `[T, H/16]`.""" g = torch.Generator(device=DEV).manual_seed(seed) x = torch.randn(tokens, H, generator=g, device=DEV).to(torch.bfloat16) x[:zero_tokens] = 0 gs = torch.tensor([1.0 / ISC1], dtype=torch.float32, device=DEV) - q, sf = torch.ops.trtllm.fp4_quantize(x, gs, VEC, False, swizzled) + q, sf = torch.ops.trtllm.fp4_quantize(x, gs, VEC, False, False) return q.view(torch.float4_e2m1fn_x2), sf.view(tokens, H // VEC) -def _ids(dist: str, tokens: int, tile: int, seed: int, num_experts: int = E) -> torch.Tensor: - """Expert ids `[tokens, K]` int32, every row naming K distinct experts.""" +def _ids(dist: str, tokens: int, seed: int) -> torch.Tensor: + """Expert ids `[tokens, K]` int32: each row's top-K of uniform scores.""" + assert dist == "router", dist g = torch.Generator(device=DEV).manual_seed(seed) - if dist == "router": - scores = torch.rand(tokens, num_experts, device=DEV, generator=g) - return scores.topk(K, dim=-1).indices.to(torch.int32) - if dist == "hot": - hot = torch.randperm(num_experts, device=DEV, generator=g)[:K] - slots = torch.rand(tokens, K, device=DEV, generator=g).argsort(dim=-1) - return hot[slots].to(torch.int32) - if dist == "tile_edges": - edge = (7, 300, 511) - scores = torch.rand(tokens, num_experts, device=DEV, generator=g) - scores[:, list(edge)] = -1.0 - ids = scores.topk(K, dim=-1).indices.to(torch.int32) - for slot, (expert, rows) in enumerate(zip(edge, (tile - 1, tile, tile + 1))): - ids[:rows, slot] = expert - return ids - if dist == "round_robin": - base = torch.arange(tokens, device=DEV).unsqueeze(1) * K - return ((base + torch.arange(K, device=DEV)) % num_experts).to(torch.int32) - raise ValueError(dist) - - -def _sort(ids: torch.Tensor, tile: int, num_experts: int = E): + scores = torch.rand(tokens, E, device=DEV, generator=g) + return scores.topk(K, dim=-1).indices.to(torch.int32) + + +def _sort(ids: torch.Tensor, tile: int): """moe_sort as the CuteDSL MoE calls it; the routing weights do not reach FC1.""" scales = torch.full(ids.shape, 0.1, dtype=torch.float32, device=DEV) return torch.ops.trtllm.moe_sort( token_selected_experts=ids, token_final_scales=scales, - num_experts=num_experts, + num_experts=E, top_k=K, local_expert_offset=0, - local_num_experts=num_experts, + local_num_experts=E, tile_tokens_dim=tile, ) -def _args(x, x_sf, sort, tile, gsf, layout="group64"): +def _args(x, x_sf, sort, tile, gsf): """The keyword arguments the target passes, by the op's own names.""" - weight, weight_scale, alpha = _weights(layout) + weight, weight_scale, alpha = _weights() t2e, t2mn, _e2p, p2e, _total, n_tiles = sort return dict( input=x, @@ -303,7 +270,7 @@ def _args(x, x_sf, sort, tile, gsf, layout="group64"): def _problem(tokens, tile, seed, dist="router", gsf="calibrated", zero_tokens=0): x, x_sf = _activation(tokens, seed, zero_tokens) - sort = _sort(_ids(dist, tokens, tile, seed + 1), tile) + sort = _sort(_ids(dist, tokens, seed + 1), tile) return _args(x, x_sf, sort, tile, GSF[gsf]), sort @@ -376,14 +343,15 @@ def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves th # ── invariance: a pair's output depends on the pair alone ───────────────── -@pytest.mark.parametrize("tokens", (128, 8192)) -def test_tiles_and_tactics_agree_per_pair(tokens, certifying): # noqa: F811 +def test_tiles_and_tactics_agree_per_pair(certifying): # noqa: F811 """Tile 128 / 256, MMA N 128 / 256 and the tuner's fallback all produce every - pair's codes and block scales bit for bit the same, so a cell certifies the - tactic the autotuner picks in serving as well as the one it ran.""" + pair's codes and block scales bit for bit the same at the 8192-row prefill, so a + cell certifies the tactic the autotuner picks in serving as well as the one it + ran.""" certifying(op) + tokens = 8192 x, x_sf = _activation(tokens, seed=tokens) - ids = _ids("router", tokens, 128, seed=tokens + 1) + ids = _ids("router", tokens, seed=tokens + 1) seen = {} for tile in (128, 256): sort = _sort(ids, tile) @@ -398,218 +366,15 @@ def test_tiles_and_tactics_agree_per_pair(tokens, certifying): # noqa: F811 _assert_pairs_equal(got, first, f"T={tokens} {key} vs (128, 128)") -def test_row_order_and_repeat_do_not_change_a_pair(certifying): # noqa: F811 - """moe_sort orders an expert's rows by atomics (from T = 5). Shuffling every - expert's run of `permuted_idx_to_expanded_idx`, and calling twice, leaves each - pair's output bitwise unchanged: the FC1 result is reproducible run to run.""" - certifying(op) - tokens, tile = 4096, 128 - args, sort = _problem(tokens, tile, seed=77) - base = _pairs(op(**args), sort, tokens) - _assert_pairs_equal(_pairs(op(**args), sort, tokens), base, "repeat call") - - t2mn, e2p, p2e, n_tiles = sort[1], sort[2], sort[3], int(sort[5]) - g = torch.Generator(device=DEV).manual_seed(5) - p2e_s = p2e.clone() - starts = torch.arange(n_tiles, device=DEV) * tile - run_of = torch.full((p2e.numel(),), -1, dtype=torch.long, device=DEV) - for i, (start, limit) in enumerate(zip(starts.tolist(), t2mn[:n_tiles].tolist())): - run_of[start:limit] = int(sort[0][i]) - for e in torch.unique(run_of[run_of >= 0]).tolist(): - rows = torch.nonzero(run_of == e).squeeze(1) - p2e_s[rows] = p2e[rows[torch.randperm(rows.numel(), generator=g, device=DEV)]] - e2p_s = e2p.clone().reshape(-1) - real = run_of >= 0 - e2p_s[p2e_s[real].long()] = torch.nonzero(real).squeeze(1).to(torch.int32) - shuffled = (sort[0], sort[1], e2p_s.view_as(e2p), p2e_s, sort[4], sort[5]) - args_s = dict(args, permuted_idx_to_expanded_idx=p2e_s) - out_s = op(**args_s) - _compare(out_s, op.reference(**args_s)) - _assert_pairs_equal(_pairs(out_s, shuffled, tokens), base, "shuffled rows") - - -def test_eager_rows_prefix_bitwise(certifying): # noqa: F811 - """The 8192-row call is gated against the reference, and every prefix call's - pairs are bitwise the 8192-row call's.""" - certifying(op) - tokens, tile = 8192, 128 - x, x_sf = _activation(tokens, seed=8) - ids = _ids("router", tokens, tile, seed=9) - sort = _sort(ids, tile) - args = _args(x, x_sf, sort, tile, GSF["calibrated"]) - full = op(**args) - _compare(full, op.reference(**args)) - full_pairs = _pairs(full, sort, tokens) - for t in _PREFIXES: - sub_sort = _sort(ids[:t].contiguous(), tile) - sub = op(**_args(x[:t], x_sf[:t], sub_sort, tile, GSF["calibrated"])) - got = _pairs(sub, sub_sort, t) - n = t * K - _assert_pairs_equal(got, (full_pairs[0][:n], full_pairs[1][:n]), f"prefix T={t}") - - -# ── CUDA graphs: captured row counts, on new data ──────────────────────── - - -def _capture(fn): - side = torch.cuda.Stream() - side.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(side): - for _ in range(2): - fn() - torch.cuda.current_stream().wait_stream(side) - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - static = fn() - return graph, static - - -def test_captured_rows_replay_on_new_data(certifying): # noqa: F811 - """moe_sort and this op are captured together, as the target's graph holds them; - each graph is replayed on two fresh activations and id sets, and its written - rows are gated against the reference and bitwise against an eager call on the - graph's own sort output. At the target's tile, 128.""" - certifying(op) - tile = 128 - for tokens in _CAPTURED_ROWS: - x, x_sf = _activation(tokens, seed=tokens) - ids = _ids("router", tokens, tile, seed=500 + tokens) - scales = torch.full(ids.shape, 0.1, dtype=torch.float32, device=DEV) - base = _args(x, x_sf, _sort(ids, tile), tile, GSF["calibrated"]) - state = {} - - def step(ids=ids, scales=scales, base=base, state=state, tile=tile): - sort = torch.ops.trtllm.moe_sort( - token_selected_experts=ids, - token_final_scales=scales, - num_experts=E, - top_k=K, - local_expert_offset=0, - local_num_experts=E, - tile_tokens_dim=tile, - ) - state["sort"] = sort - return op( - **dict( - base, - tile_idx_to_group_idx=sort[0], - tile_idx_to_mn_limit=sort[1], - permuted_idx_to_expanded_idx=sort[3], - num_non_exiting_tiles=sort[5], - ) - ) - - graph, static_out = _capture(step) - static_sort = state["sort"] - for trial in (1, 2): - nx, nx_sf = _activation(tokens, seed=1000 * trial + tokens) - x.view(torch.uint8).copy_(nx.view(torch.uint8)) - x_sf.copy_(nx_sf) - ids.copy_(_ids("router", tokens, tile, seed=7000 * trial + tokens)) - graph.replay() - args = dict( - base, - tile_idx_to_group_idx=static_sort[0], - tile_idx_to_mn_limit=static_sort[1], - permuted_idx_to_expanded_idx=static_sort[3], - num_non_exiting_tiles=static_sort[5], - ) - ref = op.reference(**args) - _compare(static_out, ref) - eager = op(**{k: _clone(v) for k, v in args.items()}) - _compare(eager, ref) - rows = ref.rows - assert torch.equal( - static_out[0].view(torch.uint8)[rows], eager[0].view(torch.uint8)[rows] - ), (tokens, trial) - m = rows.numel() - sf_g = unswizzle(static_out[1], m, INTER // VEC)[rows] - sf_e = unswizzle(eager[1], m, INTER // VEC)[rows] - assert torch.equal(sf_g, sf_e), (tokens, trial) - del graph, static_out - torch.cuda.synchronize() - - -@pytest.mark.parametrize("tokens", (3, 300)) -def test_unwritten_regions(tokens): - """A sentinel planted in a captured graph's outputs: which bytes replay writes. - The rows the sort filled hold the result; the padding rows of a written tile - are written with zero block scales and zero codes; the tiles past - `num_non_exiting_tiles` are not written -- under a graph they keep whatever - the buffer held. Tile 128.""" - tile = 128 - args, sort = _problem(tokens, tile, seed=11) - graph, static_out = _capture(lambda: op(**args)) - c, c_sf = static_out - m = c.shape[0] - for trial in (1, 2): - nx, nx_sf = _activation(tokens, seed=60 + trial) - args["input"].view(torch.uint8).copy_(nx.view(torch.uint8)) - args["input_scale"].copy_(nx_sf) - c.view(torch.uint8).fill_(SENTINEL) - c_sf.fill_(SENTINEL) - graph.replay() - ref = op.reference(**args) - _compare(static_out, ref) - n_tiles = int(sort[5]) - span = n_tiles * tile - codes = c.view(torch.uint8) - sf = unswizzle(c_sf, m, INTER // VEC) - assert bool((codes[span:] == SENTINEL).all()) and bool((sf[span:] == SENTINEL).all()), ( - "a tile past num_non_exiting_tiles was written" - ) - pad = ~ref.rows[:span] - assert int(pad.sum()) > 0 - assert bool((sf[:span][pad] == 0).all()), "a padding row's block scales are not 0" - assert bool((codes[:span][pad] == 0).all()), "a padding row's codes are not 0" - del graph - - -# ── the gate sees a wrong layout ────────────────────────────────────────── - - -@pytest.mark.parametrize( - "layout", - ("rowwise", "scales_unpermuted", "swizzled_input_scales", "alpha_rolled"), -) -def test_wrong_layouts_are_rejected(layout): - """Each operand prepared the wrong way -- the weight in trtllm-gen's row - interleave, its scales not given the weight's row permutation, the - activation's scales swizzled, alpha off by one expert -- returns finite, - plausibly scaled output that `_compare` must reject against the reference of - the correctly prepared operands.""" - tokens, tile = 128, 128 - x, x_sf = _activation(tokens, seed=3) - sort = _sort(_ids("router", tokens, tile, seed=4), tile) - good = _args(x, x_sf, sort, tile, GSF["calibrated"]) - if layout == "swizzled_input_scales": - _xs, sf_swz = _activation(tokens, seed=3, swizzled=True) - bad = dict(good, input_scale=sf_swz.reshape(tokens, H // VEC)) - elif layout == "alpha_rolled": - bad = dict(good, alpha=good["alpha"].roll(1).contiguous()) - else: - bad = _args(x, x_sf, sort, tile, GSF["calibrated"], layout=layout) - out = op(**bad) - ref = op.reference(**good) - with pytest.raises(AssertionError): - _compare(out, ref) - - -# ── guards and loud rejections ──────────────────────────────────────────── +# ── the guard ───────────────────────────────────────────────────────────── def _strided(t: torch.Tensor) -> torch.Tensor: """The same values as `t`, as a non-contiguous view into a wider buffer.""" u = t.view(torch.uint8) if t.element_size() == 1 else t - if u.dim() == 1: - wide = torch.zeros(2 * u.numel(), dtype=u.dtype, device=DEV) - wide[::2] = u - view = wide[::2] - else: - wide = torch.zeros(*u.shape[:-1], 2 * u.shape[-1], dtype=u.dtype, device=DEV) - wide[..., : u.shape[-1]] = u - view = wide[..., : u.shape[-1]] + wide = torch.zeros(*u.shape[:-1], 2 * u.shape[-1], dtype=u.dtype, device=DEV) + wide[..., : u.shape[-1]] = u + view = wide[..., : u.shape[-1]] return view.view(t.dtype) if t.element_size() == 1 else view @@ -617,12 +382,11 @@ def _u8(t: torch.Tensor) -> torch.Tensor: return t.view(torch.uint8) if t.element_size() == 1 else t -@pytest.mark.parametrize("name", ("input", "weight_scale")) -def test_noncontiguous_operand_answers_wrongly_and_is_refused(name): +def test_noncontiguous_operand_answers_wrongly_and_is_refused(): """The runner hands data_ptr() to the kernel and never reads a stride: the - same values through a strided view come back finite and wrong. One activation - and one weight operand; the guard checks all five of input, input_scale, - weight, weight_scale and alpha the same way.""" + same values through a strided view come back finite and wrong. The activation; + the guard checks input_scale, weight, weight_scale and alpha the same way.""" + name = "input" tokens, tile = 64, 128 args, _ = _problem(tokens, tile, seed=21) ref = op.reference(**args) @@ -633,66 +397,3 @@ def test_noncontiguous_operand_answers_wrongly_and_is_refused(name): _compare(op(**bad), ref) with validating(op), pytest.raises(AssertionError, match=f"^{name} must be contiguous"): op(**bad) - - -def test_short_alpha_answers_wrongly_and_is_refused(): - """A 1-element alpha (one per layer instead of one per expert): the kernel reads - alpha[expert] past its end without raising -- here into zeros of the same - allocation, so every expert but the first computes zero.""" - tokens, tile = 64, 128 - args, _ = _problem(tokens, tile, seed=31) - ref = op.reference(**args) - buf = torch.zeros(E, dtype=torch.float32, device=DEV) - buf[0] = args["alpha"][0] - bad = dict(args, alpha=buf[:1]) - with pytest.raises(AssertionError): - _compare(op(**bad), ref) - with ( - validating(op), - pytest.raises(AssertionError, match="alpha has 1 elements for 512 experts"), - ): - op(**bad) - - -def test_short_weight_scale_answers_wrongly_and_is_refused(): - """A weight_scale holding half the bytes the weight needs: the kernel reads past - its end without raising -- here into zeros of the same allocation, so the upper - half of the experts compute zero. Past the end of an allocation the same call - faults with an illegal memory access.""" - tokens, tile = 64, 128 - args, _ = _problem(tokens, tile, seed=32) - ref = op.reference(**args) - ws = args["weight_scale"].reshape(-1) - half = ws.numel() // 2 - buf = torch.zeros(ws.numel(), dtype=torch.uint8, device=DEV) - buf[:half] = ws[:half] - bad = dict(args, weight_scale=buf[:half].view(E, INTER, H // VEC)) - with pytest.raises(AssertionError): - _compare(op(**bad), ref) - with validating(op), pytest.raises(AssertionError, match="weight_scale holds"): - op(**bad) - - -def test_loud_rejections(): - """Raised before the launch: a tile size other than the sort's, either way - round; 64, which the runner does not build; a wrong scale width, a 2-element - global_sf, a bf16 alpha, scales typed e4m3 and an activation typed uint8.""" - tokens = 16 - args, _ = _problem(tokens, 128, seed=41) - with pytest.raises(AssertionError): - op(**dict(args, tile_size=256)) - args256, _ = _problem(tokens, 256, seed=42) - with pytest.raises(AssertionError): - op(**dict(args256, tile_size=128)) - with pytest.raises(ValueError, match="Tile size 64 is not supported"): - op(**dict(args, tile_size=64)) - with pytest.raises(AssertionError): - op(**dict(args, input_scale=args["input_scale"][:, :-4].contiguous())) - with pytest.raises(AssertionError): - op(**dict(args, global_sf=torch.ones(2, dtype=torch.float32, device=DEV))) - with pytest.raises(AssertionError): - op(**dict(args, alpha=args["alpha"].bfloat16())) - with pytest.raises(AssertionError): - op(**dict(args, weight_scale=args["weight_scale"].view(torch.float8_e4m3fn))) - with pytest.raises(AssertionError): - op(**dict(args, input=args["input"].view(torch.uint8))) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_grouped_gemm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_grouped_gemm.py index f86e904bc5e5..f744f13a4402 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_grouped_gemm.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_grouped_gemm.py @@ -6,10 +6,8 @@ 512 experts' down_proj in the CuteDSL weight layout, the input as the CuteDSL FC1 returns it (permuted NVFP4 rows, swizzled block scales, zero padding rows, unwritten bytes past the live tiles), the sort as moe_sort returns it -- and -holds the claims no cell carries: per-pair invariance across tiles, tactics, -row order and prefixes; CUDA-graph replay of sort, FC1 and FC2 at the captured -row counts; the regions the op leaves unwritten; that the gate rejects a -wrong layout of each operand; the guards and the loud rejections. +holds the claims no cell carries: per-pair invariance across tiles and tactics, +and the contiguity guard. """ import functools @@ -20,13 +18,11 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.cute_dsl_nvfp4_grouped_gemm import ( # noqa: E501 Fc2Expected, - unswizzle, ) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.cute_dsl_nvfp4_grouped_gemm import ( # noqa: E501 cute_dsl_nvfp4_grouped_gemm as op, ) from tensorrt_llm._torch.autotuner import AutoTuner -from tensorrt_llm._torch.utils import ActivationType __extra_import_path__ = ["../../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name @@ -40,12 +36,8 @@ ACT_AMAX = 48.0 GSF = { "calibrated": 448 * 6 / ACT_AMAX, - "saturating": 25 * 448 * 6 / ACT_AMAX, "subnormal": 2.0**-12 * 448 * 6 / ACT_AMAX, } -#: FC1's activation input scale in the chained tests: |x| = 5 maps to the e4m3 maximum. -ISC1 = 5.0 / (448 * 6) -SENTINEL = 0x5A5A #: Tie window of `_compare`, in units of the element's accumulation magnitude #: (derived in its docstring). @@ -54,13 +46,6 @@ #: the other side of a tie (4 of them when that is more). FLIP_BUDGET = 1e-5 -#: Row counts of a decode capture (requests x MTP verify width 1..4, up to 32 -#: requests): both ends of the sort's block kernel, its first cluster-kernel T and -#: the largest capture. The eager prefix lengths straddle the captured ceiling and -#: reach the 8192-row one. -_CAPTURED_ROWS = (1, 4, 5, 128) -_PREFIXES = (1, 4, 5, 127, 128, 129, 4096, 8191) - # ── the gate ────────────────────────────────────────────────────────────── @@ -73,9 +58,8 @@ def _bf16_order(x: torch.Tensor) -> torch.Tensor: return torch.where(bits >= 0x8000, -mag, mag) -def _deviation(out: torch.Tensor, ref: Fc2Expected, rows: torch.Tensor | None = None) -> dict: - """Every disagreement between the op and the reference on the rows it writes - (or on the subset `rows` of them). +def _deviation(out: torch.Tensor, ref: Fc2Expected) -> dict: + """Every disagreement between the op and the reference on the rows it writes. An element is `equal` when its value matches the reference's `out` (+0 and -0 alike). Otherwise it is a `flip` when the two are neighbouring bf16 @@ -83,10 +67,8 @@ def _deviation(out: torch.Tensor, ref: Fc2Expected, rows: torch.Tensor | None = between them, and `outside` when it is anything else. `far` counts the disagreements that are not between neighbours at all. """ - sel = ref.rows if rows is None else rows - assert not bool((sel & ~ref.rows).any()), "deviation reads written rows only" span = ref.exact.shape[0] - idx = torch.nonzero(sel[:span]).squeeze(1) + idx = torch.nonzero(ref.rows[:span]).squeeze(1) o = out[idx] r = ref.out[idx] v = ref.exact[idx] @@ -174,22 +156,12 @@ def _down_proj(): return _quantized_experts(E, H, INTER, seed=20261013) -def _shuffle32(rows: int) -> torch.Tensor: - """trtllm-gen's 32-row block shuffle: destination row `8v + u` <- source `4u + v`.""" - idx = torch.arange(rows, device=DEV).view(rows // 32, 8, 4) - return idx.transpose(1, 2).reshape(-1) - - @functools.lru_cache(maxsize=None) -def _weights(gsf: str = "calibrated", layout: str = "cutedsl"): +def _weights(gsf: str = "calibrated"): """`(weight, weight_scale, alpha)` in the layout CuteDSL reads -- what the target's - load-time transform produces -- or one of the wrong layouts the discrimination - test feeds it. `alpha = down_proj.input_scale * weight_scale_2`, with the - input scale the inverse of FC1's `global_sf`.""" + load-time transform produces. `alpha = down_proj.input_scale * weight_scale_2`, + with the input scale the inverse of FC1's `global_sf`.""" w, s, ws2 = _down_proj() - if layout == "shuffle32": - perm = _shuffle32(H) - w, s = w[:, perm], s[:, perm] s_k = torch.ops.trtllm.block_scale_interleave(s.contiguous()).view(s.shape) alpha = (ws2 / GSF[gsf]).contiguous() return w.contiguous().view(torch.float4_e2m1fn_x2), s_k, alpha @@ -226,28 +198,12 @@ def _fc2_input(y: torch.Tensor, sort, gsf: str, garbage_seed: int = 0): return q.view(torch.float4_e2m1fn_x2), _swizzle(sf).contiguous() -def _ids(dist: str, tokens: int, tile: int, seed: int) -> torch.Tensor: - """Expert ids `[tokens, K]` int32, every row naming K distinct experts.""" +def _ids(dist: str, tokens: int, seed: int) -> torch.Tensor: + """Expert ids `[tokens, K]` int32: each row's top-K of uniform scores.""" + assert dist == "router", dist g = torch.Generator(device=DEV).manual_seed(seed) - if dist == "router": - scores = torch.rand(tokens, E, device=DEV, generator=g) - return scores.topk(K, dim=-1).indices.to(torch.int32) - if dist == "hot": - hot = torch.randperm(E, device=DEV, generator=g)[:K] - slots = torch.rand(tokens, K, device=DEV, generator=g).argsort(dim=-1) - return hot[slots].to(torch.int32) - if dist == "tile_edges": - edge = (7, 300, 511) - scores = torch.rand(tokens, E, device=DEV, generator=g) - scores[:, list(edge)] = -1.0 - ids = scores.topk(K, dim=-1).indices.to(torch.int32) - for slot, (expert, rows) in enumerate(zip(edge, (tile - 1, tile, tile + 1))): - ids[:rows, slot] = expert - return ids - if dist == "round_robin": - base = torch.arange(tokens, device=DEV).unsqueeze(1) * K - return ((base + torch.arange(K, device=DEV)) % E).to(torch.int32) - raise ValueError(dist) + scores = torch.rand(tokens, E, device=DEV, generator=g) + return scores.topk(K, dim=-1).indices.to(torch.int32) def _sort(ids: torch.Tensor, tile: int): @@ -264,9 +220,9 @@ def _sort(ids: torch.Tensor, tile: int): ) -def _args(x, x_sf, sort, tile, gsf="calibrated", layout="cutedsl"): +def _args(x, x_sf, sort, tile, gsf="calibrated"): """The keyword arguments the target passes, by the op's own names.""" - weight, weight_scale, alpha = _weights(gsf, layout) + weight, weight_scale, alpha = _weights(gsf) return dict( input=x, weight=weight, @@ -286,7 +242,7 @@ def _args(x, x_sf, sort, tile, gsf="calibrated", layout="cutedsl"): def _problem(tokens, tile, seed, dist="router", gsf="calibrated", zero_tokens=0): - sort = _sort(_ids(dist, tokens, tile, seed + 1), tile) + sort = _sort(_ids(dist, tokens, seed + 1), tile) y = _pair_rows(tokens * K, seed, zero_pairs=zero_tokens * K) x, x_sf = _fc2_input(y, sort, gsf, garbage_seed=seed + 2) return _args(x, x_sf, sort, tile, gsf), sort @@ -315,12 +271,6 @@ def _pairs(out, sort, tokens): return out[sort[2].reshape(-1).long()[: tokens * K]] -def _pair_mask(sort, m, tokens): - mask = torch.zeros(m, dtype=torch.bool, device=DEV) - mask[sort[2].reshape(-1).long()[: tokens * K]] = True - return mask - - def _clone(v): """A copy that does not depend on copy kernels for the fp4 dtype.""" if not isinstance(v, torch.Tensor): @@ -367,13 +317,14 @@ def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves th # ── invariance: a pair's row depends on the pair alone ──────────────────── -@pytest.mark.parametrize("tokens", (128, 8192)) -def test_tiles_and_tactics_agree_per_pair(tokens, certifying): # noqa: F811 +def test_tiles_and_tactics_agree_per_pair(certifying): # noqa: F811 """Tile 128 / 256, MMA N 128 / 256, cluster N 1 / 2 and the tuner's fallback - all produce every pair's row bit for bit the same, so a cell certifies the - tactic the autotuner picks in serving as well as the one it ran.""" + all produce every pair's row bit for bit the same at the 8192-row prefill, so a + cell certifies the tactic the autotuner picks in serving as well as the one it + ran.""" certifying(op) - ids = _ids("router", tokens, 128, seed=tokens + 1) + tokens = 8192 + ids = _ids("router", tokens, seed=tokens + 1) y = _pair_rows(tokens * K, seed=tokens) seen = {} for tile in (128, 256): @@ -391,260 +342,7 @@ def test_tiles_and_tactics_agree_per_pair(tokens, certifying): # noqa: F811 _assert_pairs_equal(got, first, f"T={tokens} {key} vs (128, 128, 1)") -def test_row_order_and_repeat_do_not_change_a_pair(certifying): # noqa: F811 - """moe_sort orders an expert's rows by atomics (from T = 5). Moving every pair to - another row of its expert's run, and calling twice, leaves each pair's row - bitwise unchanged: the FC2 result is reproducible run to run.""" - certifying(op) - tokens, tile = 4096, 128 - ids = _ids("router", tokens, tile, seed=78) - sort = _sort(ids, tile) - y = _pair_rows(tokens * K, seed=77) - x, x_sf = _fc2_input(y, sort, "calibrated") - args = _args(x, x_sf, sort, tile) - base = _pairs(op(**args), sort, tokens) - _assert_pairs_equal(_pairs(op(**args), sort, tokens), base, "repeat call") - - t2e, t2mn, e2p, p2e, n_tiles = sort[0], sort[1], sort[2], sort[3], int(sort[5]) - g = torch.Generator(device=DEV).manual_seed(5) - p2e_s = p2e.clone() - starts = torch.arange(n_tiles, device=DEV) * tile - run_of = torch.full((p2e.numel(),), -1, dtype=torch.long, device=DEV) - for i, (start, limit) in enumerate(zip(starts.tolist(), t2mn[:n_tiles].tolist())): - run_of[start:limit] = int(t2e[i]) - for e in torch.unique(run_of[run_of >= 0]).tolist(): - rows = torch.nonzero(run_of == e).squeeze(1) - p2e_s[rows] = p2e[rows[torch.randperm(rows.numel(), generator=g, device=DEV)]] - e2p_s = e2p.clone().reshape(-1) - real = run_of >= 0 - e2p_s[p2e_s[real].long()] = torch.nonzero(real).squeeze(1).to(torch.int32) - shuffled = (t2e, t2mn, e2p_s.view_as(e2p), p2e_s, sort[4], sort[5]) - assert not torch.equal(e2p_s.view_as(e2p), e2p), "the shuffle moved no pair" - xs, xs_sf = _fc2_input(y, shuffled, "calibrated") - args_s = _args(xs, xs_sf, shuffled, tile) - out_s = op(**args_s) - _compare(out_s, op.reference(**args_s)) - _assert_pairs_equal(_pairs(out_s, shuffled, tokens), base, "pairs moved within their runs") - - -def test_eager_rows_prefix_bitwise(certifying): # noqa: F811 - """The 8192-row call is gated against the reference, and every prefix call's - pairs are bitwise the 8192-row call's.""" - certifying(op) - tokens, tile = 8192, 128 - ids = _ids("router", tokens, tile, seed=9) - y = _pair_rows(tokens * K, seed=8) - sort = _sort(ids, tile) - x, x_sf = _fc2_input(y, sort, "calibrated") - args = _args(x, x_sf, sort, tile) - full = op(**args) - _compare(full, op.reference(**args)) - full_pairs = _pairs(full, sort, tokens) - for t in _PREFIXES: - sub_sort = _sort(ids[:t].contiguous(), tile) - xs, xs_sf = _fc2_input(y[: t * K], sub_sort, "calibrated", garbage_seed=t) - got = _pairs(op(**_args(xs, xs_sf, sub_sort, tile)), sub_sort, t) - _assert_pairs_equal(got, full_pairs[: t * K], f"prefix T={t}") - - -# ── chained from the CuteDSL FC1, eager and captured ────────────────────── - - -@functools.lru_cache(maxsize=None) -def _fc1_weights(): - """FC1 operands for the chain: `[E, 2I, H/2]` codes, swizzled scales, per-expert - alpha and the layer's `global_sf`. Their layout is FC1's business; what FC2 - sees is FC1's output, whatever the weights.""" - w, s, ws2 = _quantized_experts(E, 2 * INTER, H, seed=20261014) - s_k = torch.ops.trtllm.block_scale_interleave(s.contiguous()).view(s.shape) - alpha = (ISC1 * ws2).contiguous() - gsf = torch.tensor([GSF["calibrated"]], dtype=torch.float32, device=DEV) - return w.view(torch.float4_e2m1fn_x2), s_k, alpha, gsf - - -def _activation(tokens: int, seed: int): - """The MoE input as the target quantizes it: `fp4_quantize(x, 1/isc1, 16, False, - False)`, scales viewed `[T, H/16]`.""" - g = torch.Generator(device=DEV).manual_seed(seed) - x = torch.randn(tokens, H, generator=g, device=DEV).to(torch.bfloat16) - gs = torch.tensor([1.0 / ISC1], dtype=torch.float32, device=DEV) - q, sf = torch.ops.trtllm.fp4_quantize(x, gs, VEC, False, False) - return q.view(torch.float4_e2m1fn_x2), sf.view(tokens, H // VEC) - - -def _fc1(xq, xsf, sort, tile): - """The CuteDSL FC1, called as the target calls it; returns `(c, c_sf)`.""" - w1, s1, alpha1, gsf = _fc1_weights() - return torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell( - input=xq, - weight=w1, - input_scale=xsf, - weight_scale=s1, - alpha=alpha1, - tile_idx_to_group_idx=sort[0], - tile_idx_to_mn_limit=sort[1], - permuted_idx_to_expanded_idx=sort[3], - num_non_exiting_tiles=sort[5], - global_sf=gsf, - num_experts=E, - top_k=K, - num_local_experts=E, - local_expert_offset=0, - tile_size=tile, - scaling_vector_size=VEC, - activation_type=int(ActivationType.Swiglu), - ) - - -def test_fc1_output_feeds_it_at_prefill(certifying): # noqa: F811 - """The real FC1 output at 8192 tokens -- the target's prefill -- passes `_compare`, - and the padding rows FC1 zeroes come out zero.""" - certifying(op) - tokens, tile = 8192, 128 - xq, xsf = _activation(tokens, seed=91) - sort = _sort(_ids("router", tokens, tile, seed=92), tile) - h, h_sf = _fc1(xq, xsf, sort, tile) - args = _args(h, h_sf, sort, tile) - out = op(**args) - ref = op.reference(**args) - _compare(out, ref) - pad = ref.rows & ~_pair_mask(sort, out.shape[0], tokens) - assert int(pad.sum()) > 0 - assert bool((out[pad] == 0).all()), "a padding row of FC1's output did not come out zero" - - -def _capture(fn): - side = torch.cuda.Stream() - side.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(side): - for _ in range(2): - fn() - torch.cuda.current_stream().wait_stream(side) - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - static = fn() - return graph, static - - -def test_captured_rows_replay_on_new_data(certifying): # noqa: F811 - """moe_sort, FC1 and this op are captured together, as the target's graph holds - them; each graph is replayed on two fresh activations and id sets, and its - written rows are gated against the reference of the graph's own FC1 output and - sort, and bitwise against an eager call on them. At the target's tile, 128.""" - certifying(op) - tile = 128 - for tokens in _CAPTURED_ROWS: - xq, xsf = _activation(tokens, seed=tokens) - ids = _ids("router", tokens, tile, seed=500 + tokens) - scales = torch.full(ids.shape, 0.1, dtype=torch.float32, device=DEV) - state = {} - - def step(xq=xq, xsf=xsf, ids=ids, scales=scales, state=state, tile=tile): - sort = torch.ops.trtllm.moe_sort( - token_selected_experts=ids, - token_final_scales=scales, - num_experts=E, - top_k=K, - local_expert_offset=0, - local_num_experts=E, - tile_tokens_dim=tile, - ) - h, h_sf = _fc1(xq, xsf, sort, tile) - state["sort"], state["h"], state["h_sf"] = sort, h, h_sf - return op(**_args(h, h_sf, sort, tile)) - - graph, static_out = _capture(step) - for trial in (1, 2): - nx, nx_sf = _activation(tokens, seed=1000 * trial + tokens) - xq.view(torch.uint8).copy_(nx.view(torch.uint8)) - xsf.copy_(nx_sf) - ids.copy_(_ids("router", tokens, tile, seed=7000 * trial + tokens)) - graph.replay() - args = _args(state["h"], state["h_sf"], state["sort"], tile) - ref = op.reference(**args) - _compare(static_out, ref) - eager = op(**{k: _clone(v) for k, v in args.items()}) - rows = ref.rows - assert torch.equal(_bits(static_out[rows]), _bits(eager[rows])), (tokens, trial) - del graph, static_out - torch.cuda.synchronize() - - -# ── what the op writes ──────────────────────────────────────────────────── - - -@pytest.mark.parametrize("tokens", (3, 300)) -def test_unwritten_regions(tokens): - """A sentinel planted in a captured graph's output: which rows replay writes. - Every row of the live tiles is written -- the padding rows computed from - whatever the input holds there, zero for FC1's zeroed padding -- and the rows - past `num_non_exiting_tiles` are not: under a graph they keep whatever the - buffer held. Tile 128.""" - tile = 128 - args, sort = _problem(tokens, tile, seed=11) - graph, out = _capture(lambda: op(**args)) - m = out.shape[0] - span = int(sort[5]) * tile - pairs = _pair_mask(sort, m, tokens) - pad = torch.zeros(m, dtype=torch.bool, device=DEV) - pad[:span] = ~pairs[:span] - assert int(pad.sum()) > 0 - # With no more pairs than experts the sort sizes one tile per pair; when no - # two pairs share an expert (T = 3 here) every tile is live. - assert span < m or tokens * K <= E, "the sort left no tile unused" - _bits(out).fill_(SENTINEL) - graph.replay() - _compare(out, op.reference(**args)) - untouched = (_bits(out) == SENTINEL).all(dim=1) - assert bool(untouched[span:].all()), "a row past num_non_exiting_tiles was written" - assert not bool(untouched[:span].any()), "a row of a live tile was not written" - assert bool((out[pad] == 0).all()), "a zero padding row did not come out zero" - pair_rows = out[pairs].clone() - - # Nonzero codes (all 0.5) under unit block scales in the padding rows. - codes = args["input"].view(torch.uint8) - codes[pad] = 0x11 - sf = unswizzle(args["input_scale"], m, NB).clone() - sf[pad] = 0x38 - args["input_scale"].copy_(_swizzle(sf)) - graph.replay() - assert bool((out[pad] != 0).any()), "nonzero padding input still gave zero rows" - assert torch.equal(_bits(out[pairs]), _bits(pair_rows)), "padding input changed a pair row" - _compare(out, op.reference(**args)) - del graph - - -# ── the gate sees a wrong layout ────────────────────────────────────────── - - -@pytest.mark.parametrize( - "layout", - ("shuffle32", "linear_input_scales", "alpha_per_layer"), -) -def test_wrong_layouts_are_rejected(layout): - """Each operand prepared the wrong way -- the weight in trtllm-gen's 32-row - shuffle, the input's scales linear, one alpha for the whole layer -- returns - finite, plausibly scaled output that `_compare` must reject against the - reference of the correctly prepared operands.""" - tokens, tile = 128, 128 - good, sort = _problem(tokens, tile, seed=3) - m = good["input"].shape[0] - if layout == "shuffle32": - bad = _args(good["input"], good["input_scale"], sort, tile, layout=layout) - elif layout == "linear_input_scales": - linear = unswizzle(good["input_scale"], m, NB).reshape(-1).contiguous() - bad = dict(good, input_scale=linear) - else: - bad = dict(good, alpha=torch.full_like(good["alpha"], float(good["alpha"].mean()))) - out = op(**bad) - ref = op.reference(**good) - assert _deviation(out, ref, rows=_pair_mask(sort, m, tokens))["nonfinite"] == 0 - with pytest.raises(AssertionError): - _compare(out, ref) - - -# ── guards and loud rejections ──────────────────────────────────────────── +# ── the guard ───────────────────────────────────────────────────────────── def _u8(t: torch.Tensor) -> torch.Tensor: @@ -654,23 +352,17 @@ def _u8(t: torch.Tensor) -> torch.Tensor: def _strided(t: torch.Tensor) -> torch.Tensor: """The same values as `t`, as a non-contiguous view into a wider buffer.""" u = _u8(t) - if u.dim() == 1: - wide = torch.zeros(2 * u.numel(), dtype=u.dtype, device=DEV) - wide[::2] = u - view = wide[::2] - else: - wide = torch.zeros(*u.shape[:-1], 2 * u.shape[-1], dtype=u.dtype, device=DEV) - wide[..., : u.shape[-1]] = u - view = wide[..., : u.shape[-1]] + wide = torch.zeros(*u.shape[:-1], 2 * u.shape[-1], dtype=u.dtype, device=DEV) + wide[..., : u.shape[-1]] = u + view = wide[..., : u.shape[-1]] return view.view(t.dtype) if t.element_size() == 1 else view -@pytest.mark.parametrize("name", ("input", "weight_scale", "tile_idx_to_group_idx")) -def test_noncontiguous_operand_answers_wrongly_and_is_refused(name): +def test_noncontiguous_operand_answers_wrongly_and_is_refused(): """The runner hands data_ptr() to the kernel and never reads a stride: the - same values through a strided view come back finite and wrong. One activation, - one weight and one sort operand; the guard checks input_scale, weight and - alpha the same way.""" + same values through a strided view come back finite and wrong. The activation; + the guard checks the other five operands the same way.""" + name = "input" tokens, tile = 64, 128 args, _ = _problem(tokens, tile, seed=21) ref = op.reference(**args) @@ -681,38 +373,3 @@ def test_noncontiguous_operand_answers_wrongly_and_is_refused(name): _compare(op(**bad), ref) with validating(op), pytest.raises(AssertionError, match=f"^{name} must be contiguous"): op(**bad) - - -def test_loud_rejections(): - """Raised before the launch: a tile size other than the sort's, either way - round; 64, which the runner does not build; a non-bf16 output; and a wrong - dtype, rank or size of every tensor operand, each tripping a runner assert.""" - tokens = 16 - args, _ = _problem(tokens, 128, seed=41) - m = args["input"].shape[0] - with pytest.raises(AssertionError): - op(**dict(args, tile_size=256)) - args256, _ = _problem(tokens, 256, seed=42) - with pytest.raises(AssertionError): - op(**dict(args256, tile_size=128)) - with pytest.raises(ValueError, match="supports tile_size .* 128 and 256 only, but got 64"): - op(**dict(args, tile_size=64)) - rejected = dict( - output_dtype=torch.float16, - alpha=args["alpha"].bfloat16(), - input=args["input"].view(torch.uint8), - input_scale=args["input_scale"].view(m, NB), - weight_scale=args["weight_scale"].view(torch.float8_e4m3fn), - scaling_vector_size=32, - tile_idx_to_group_idx=args["tile_idx_to_group_idx"].long(), - num_non_exiting_tiles=torch.zeros(2, dtype=torch.int32, device=DEV), - ) - short = dict( - alpha=args["alpha"][:1].contiguous(), - weight_scale=args["weight_scale"][: E // 2], - input_scale=args["input_scale"][:-4].contiguous(), - tile_idx_to_group_idx=args["tile_idx_to_group_idx"][:-1].contiguous(), - ) - for name, value in list(rejected.items()) + list(short.items()): - with pytest.raises(AssertionError): - op(**dict(args, **{name: value})) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_weight_path.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_weight_path.py index 08e0c2758cd0..ddef4b35ea36 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_weight_path.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_weight_path.py @@ -13,11 +13,12 @@ Two routings hit all 512 experts. In the first every token's only non-zero routing weight is on its own expert, so each expert's output is checked alone: a layout error confined to one -expert (a permutation, a scale swizzle, an alpha) cannot hide in a sum over ten. The second is -a weighted top-10. Both backends are held to a float64 reference over the dequantized -checkpoint tensors, with `down_proj.input_scale` spread across the experts as a calibrated -checkpoint's is. With a uniform one the two backends requantize the SwiGLU output -identically, so on the weighted routing they are also held to each other. +expert (a permutation, a scale swizzle, an alpha) cannot hide in a sum over ten. On it both +backends are held to a float64 reference over the dequantized checkpoint tensors, with +`down_proj.input_scale` spread across the experts as a calibrated checkpoint's is. The second +is a weighted top-10: with a uniform `down_proj.input_scale` the two backends requantize the +SwiGLU output identically, so on it they are held to each other -- the area's one +cross-backend check. The reference does not model the kernels' accumulation. It models the one rounding step that is not exact: the SwiGLU output is requantized to NVFP4 before FC2, at the backend's @@ -66,7 +67,6 @@ WEIGHTED_TOKENS = 1024 BACKENDS = ("CUTEDSL", "TRTLLM") -CASES = ("one_hot", "weighted") #: The float64 reference's gate, relative per token. Both backends sit at 2.5e-3 or below: #: the bf16 rounding of the output is 1e-3 of it, the rest is the reference's fp32 / tie #: differences. A wrong layout is ~1; a wrong per-expert scalar is its own relative error. @@ -348,22 +348,20 @@ def _token_error(out: torch.Tensor, ref: torch.Tensor) -> torch.Tensor: # ── tests ───────────────────────────────────────────────────────────────── -def test_routings_hit_every_expert() -> None: - for case in CASES: - _, ids, weights = _problem(case) - assert torch.unique(ids[weights != 0]).numel() == E, case - - -@pytest.mark.parametrize("case", CASES) @pytest.mark.parametrize("backend", BACKENDS) -def test_routed_output_matches_the_float64_reference(backend: str, case: str) -> None: +def test_routed_output_matches_the_float64_reference(backend: str) -> None: + case = "one_hot" + _, ids, weights = _problem(case) + assert torch.unique(ids[weights != 0]).numel() == E, "the routing misses an expert" err = _token_error(_routed(backend, case, True), _reference(backend, case, True)) print(f"[err] {backend} {case}: mean {err.mean():.4e} max {err.max():.4e}") assert err.max() <= REFERENCE_TOLERANCE, (err.mean().item(), err.max().item()) -@pytest.mark.parametrize("case", ("weighted",)) -def test_cutedsl_matches_trtllm_gen_under_a_uniform_input_scale(case: str) -> None: +def test_cutedsl_matches_trtllm_gen_under_a_uniform_input_scale() -> None: + case = "weighted" + _, ids, weights = _problem(case) + assert torch.unique(ids[weights != 0]).numel() == E, "the routing misses an expert" err = _token_error(_routed("CUTEDSL", case, False), _routed("TRTLLM", case, False)) print(f"[err] CUTEDSL vs TRTLLM {case}: mean {err.mean():.4e} max {err.max():.4e}") assert err.max() <= BACKEND_TOLERANCE, (err.mean().item(), err.max().item()) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py index 2fdc11d84ad4..c0ead50efd89 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py @@ -7,7 +7,8 @@ and it sits inside the pre-routed entry's band against this entry's own fp32 reference. The NVFP4 operand construction (layout, activations, scalars, the band) is the pre-routed entry's test's, loaded by path so the two tests -cannot drift apart on what a valid operand is. +cannot drift apart on what a valid operand is. Beyond the cells: the refusal +of an uncertified (num_experts, top_k) pair. """ import importlib.util @@ -118,32 +119,6 @@ def test_certified_cells(cell, certifying): # noqa: F811 H._compare(y, exp.to(torch.bfloat16)) -def test_coarse_cell_exercises_tie_breaking(): - """The coarse cell is only worth its slot if ties decide real slots: the - opposite tie order (ties to the HIGHER expert index) must pick a different - expert set on most rows, so the cell's bitwise check pins the direction.""" - spec = next(c.spec for c in op.CELLS if c.spec["dist"] == "coarse") - e, k = spec["num_experts"], spec["top_k"] - gen = torch.Generator(device=DEV).manual_seed(7) - logits = _logits(spec, gen) - ids, _ = renorm_moe_routing_op(logits, k, torch.bfloat16) - flipped = torch.flip(logits, [-1]) - order = torch.sort(flipped, dim=-1, descending=True, stable=True).indices[:, :k] - high = (e - 1 - order).int() - differ = (ids.sort(-1).values != high.sort(-1).values).any(-1).float().mean().item() - assert differ > 0.5, f"only {differ:.2f} of rows break ties differently" - - -def test_strided_logits_refused(): - # More than one row: a one-row column slice is contiguous by definition. - spec = next(c.spec for c in op.CELLS if c.spec["tokens"] == 128) - args, _, data, sf, _, logits, _, _, _, scalars = _operands(spec, 11) - wide = torch.zeros(spec["tokens"], spec["num_experts"] + 8, device=DEV, dtype=logits.dtype) - wide[:, : spec["num_experts"]] = logits - with validating(op), pytest.raises(AssertionError, match="routing_logits must be contiguous"): - _routed(args, data, sf, wide[:, : spec["num_experts"]], scalars, spec) - - def test_uncertified_pair_refused(): spec = op.CELLS[0].spec assert (spec["num_experts"], 9) not in CERTIFIED_PAIRS diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py index f1df7ac292a4..5f2aeabd5cf1 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py @@ -6,8 +6,8 @@ power-of-two scales, so every weight is exact); this entry's layout is built from their raw per-expert tensors by the target's own `weights.py` transforms, so the test also pins those. Per cell: inside the routed entry's band against -the fp32 reference, and close to the trtllm-gen routed entry on the same -experts. +the fp32 reference. Beyond the cells: the guard against a per-expert FC2 input +scale. """ import importlib.util @@ -20,9 +20,6 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 weights as W, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp4_block_scale_moe_renorm_routed import ( # noqa: E501 - fp4_block_scale_moe_renorm_routed as routed, -) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp4_moe_renorm_cute_dsl import ( # noqa: E501 fp4_moe_renorm_cute_dsl as op, ) @@ -113,35 +110,6 @@ def test_certified_cells(cell, certifying): # noqa: F811 H._compare(y, exp.to(torch.bfloat16)) -def test_close_to_the_trtllm_gen_entry(): - """Same experts, same routing: within the routed entry's own band of each - other's output (the two differ only in accumulation order). At the row count - where the target switches from one to the other.""" - tokens = 4096 - spec = {**op.CELLS[0].spec, "tokens": tokens} - args, ref, data, sf, x, logits, g2 = _operands(spec, 4000 + tokens) - y = _cute_dsl(ref, data, sf, logits, g2, spec) - s1, sg, s2 = H._scalars(spec["num_experts"], G1, g2) - y_gen = routed( - logits, - data, - sf.view(torch.float8_e4m3fn), - args["gemm1_weights"], - args["gemm1_weights_scale"], - args["gemm2_weights"], - args["gemm2_weights_scale"], - s1, - sg, - s2, - spec["num_experts"], - spec["top_k"], - spec["intermediate"], - ) - rel = ((y.float() - y_gen.float()).norm() / y_gen.float().norm()).item() - print(f"\n {tokens} rows: relative RMS difference from trtllm-gen {rel:.2e}") - assert rel < 1e-2 - - def test_guard_refuses_a_per_expert_fc2_scale(): spec = {**op.CELLS[0].spec, "tokens": 1024} args, ref, data, sf, x, logits, g2 = _operands(spec, 9) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp8_block_scale_moe_runner.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp8_block_scale_moe_runner.py index d463c4e26622..b2d77dcf4336 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp8_block_scale_moe_runner.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp8_block_scale_moe_runner.py @@ -3,9 +3,9 @@ """GPU test for the fp8_block_scale_moe_runner catalog entry. The entry owns the contract: `CELLS`, `reference`, `is_valid`. This -file turns a cell into operands and holds the claims no cell carries: the -intermediate requantization pinned bit-exactly, how far each wrong reference -lands, the expert-window bookkeeping, and the guards. +file turns a cell into operands and holds the claims no cell carries: a warm +autotuned tactic is bitwise the cold fallback the cells run, and the +contiguity guard. """ import zlib @@ -13,24 +13,15 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import true_fp32_matmul from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp8_block_scale_moe_runner import ( # noqa: E501 - ACT_AMAX_FLOOR, BLOCK, - E4M3_MAX, - dequantize_block_weight, - requantize_1x128, ) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp8_block_scale_moe_runner import ( # noqa: E501 fp8_block_scale_moe_runner as op, ) __extra_import_path__ = ["../../../.."] -from _validating import ( # noqa: E402,F401 — path above; fixture by name - assert_within_ulp, - certifying, - validating, -) +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "fp8_block_scale_moe_runner requires a CUDA device" @@ -53,8 +44,7 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: 0.342. Wrong references at the checkpoint geometry, RMS: gate/up swap 65, transposed activation scales 43, strided hidden_states 94, no FC1-output requant 3.5, neither requant 3.1, no activation requant 1.58 - -- the closest, which is why the intermediate requantization is also - pinned bit-exactly by the test. + -- the closest. """ assert out.dtype == ref.dtype == torch.bfloat16, (out.dtype, ref.dtype) assert out.shape == ref.shape, (out.shape, ref.shape) @@ -101,25 +91,11 @@ def _weights(hidden: int, inter: int, experts: int, seed: int): def _args( - tokens, - hidden, - inter, - num_experts, - top_k, - local_expert_offset=0, - local_num_experts=None, - gemm1_clamp_limit=None, - seed=0, + tokens, hidden, inter, num_experts, top_k, local_expert_offset=0, local_num_experts=None, seed=0 ): - """The op's arguments by name, pre-routed, for one configuration. - - Weights are drawn for the whole routing space and the window sliced out - of them, so a window and the full stack hold the same experts. - """ - local_num_experts = num_experts if local_num_experts is None else local_num_experts - stack = max(num_experts, local_expert_offset + local_num_experts) - w1, w1s, w2, w2s = _weights(hidden, inter, stack, seed=hidden + inter + stack) - window = slice(local_expert_offset, local_expert_offset + local_num_experts) + """The op's arguments by name, pre-routed, every expert local.""" + assert local_expert_offset == 0 and local_num_experts in (None, num_experts) + w1, w1s, w2, w2s = _weights(hidden, inter, num_experts, seed=hidden + inter + num_experts) g = torch.Generator(device=DEV).manual_seed(seed) hs = (2.0 * torch.randn(tokens, hidden, device=DEV, generator=g)).to(torch.float8_e4m3fn) # [H/128, T]: the transposed layout fp8_quantize_1x128 returns here. @@ -132,23 +108,23 @@ def _args( routing_bias=None, hidden_states=hs, hidden_states_scale=hs_scale, - gemm1_weights=w1[window].contiguous(), - gemm1_weights_scale=w1s[window].contiguous(), - gemm2_weights=w2[window].contiguous(), - gemm2_weights_scale=w2s[window].contiguous(), + gemm1_weights=w1, + gemm1_weights_scale=w1s, + gemm2_weights=w2, + gemm2_weights_scale=w2s, num_experts=num_experts, top_k=top_k, num_fused_shared_experts=None, n_group=None, topk_group=None, intermediate_size=inter, - local_expert_offset=local_expert_offset, - local_num_experts=local_num_experts, + local_expert_offset=0, + local_num_experts=num_experts, routed_scaling_factor=None, routing_method_type=1, topk_weights=weights.to(torch.bfloat16).contiguous(), topk_ids=ids.contiguous(), - gemm1_clamp_limit=gemm1_clamp_limit, + gemm1_clamp_limit=None, ) @@ -159,48 +135,6 @@ def _ulp(out: torch.Tensor, ref: torch.Tensor): return d.abs().max().item(), d.pow(2).mean().sqrt().item() -def _variant( - a, swap=False, quant_fc1_out=True, quant_act=True, count_repeats=True, requant=requantize_1x128 -): - """`reference` with one step changed -- the wrong readings it must refute.""" - tokens, hidden = a["hidden_states"].shape - inter, nb = a["intermediate_size"], hidden // BLOCK - act_in = ( - a["hidden_states"].float().reshape(tokens, nb, BLOCK) - * a["hidden_states_scale"].t().reshape(tokens, nb, 1) - ).reshape(tokens, hidden) - ids = a["topk_ids"].long() - local = ids - a["local_expert_offset"] - live = (local >= 0) & (local < a["local_num_experts"]) - if not count_repeats: - first = torch.ones_like(live) - for j in range(1, ids.shape[1]): - first[:, j] = (ids[:, j : j + 1] != ids[:, :j]).all(dim=1) - live &= first - out = torch.zeros(tokens, hidden, device=DEV) - with true_fp32_matmul(): - for e in torch.unique(local[live]).tolist(): - tok, slot = ((local == e) & live).nonzero(as_tuple=True) - y1 = ( - act_in[tok] - @ dequantize_block_weight(a["gemm1_weights"][e], a["gemm1_weights_scale"][e]).t() - ) - if quant_fc1_out: - y1 = requant(y1) - up, gate = y1[:, :inter], y1[:, inter:] - if swap: - up, gate = gate, up - act = torch.nn.functional.silu(gate) * up - if quant_act: - act = requant(act) - y = ( - act - @ dequantize_block_weight(a["gemm2_weights"][e], a["gemm2_weights_scale"][e]).t() - ) - out.index_add_(0, tok, y * a["topk_weights"][tok, slot].float().unsqueeze(1)) - return out.to(torch.bfloat16) - - def _seed(cell) -> int: return zlib.crc32(cell.why.encode()) @@ -220,162 +154,7 @@ def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves th _compare(out, ref) -# ── the two internal requantizations, read out bit-exactly ──────────────── - - -def _pow2_requantize(x: torch.Tensor) -> torch.Tensor: - """The MX-style round-up power-of-two block scale -- a wrong recipe here.""" - shp = x.shape - xb = x.reshape(*shp[:-1], shp[-1] // BLOCK, BLOCK) - amax = xb.abs().amax(dim=-1, keepdim=True).clamp(min=ACT_AMAX_FLOOR) - scale = torch.exp2(torch.ceil(torch.log2(amax / E4M3_MAX))) - coded = (xb / scale).clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn) - return (coded.float() * scale).reshape(shp) - - -@pytest.mark.parametrize("tokens,width", [(64, 640), (256, 128)]) -def test_intermediate_is_fp8_block_quantized(tokens, width): - """Identity down projection, one slot of weight 1.0: the output is the - requantized intermediate, one bf16 round away. The recipe reproduces - >=99.9% of it bit-exactly (measured 99.95-100%); every neighbouring recipe - reproduces under 10%. The misses are FC1 sums on an e4m3 rounding boundary. - """ - num_experts = 8 - a = _args(tokens, width, width, num_experts, 1, seed=width + tokens) - eye = torch.eye(width, device=DEV).to(torch.float8_e4m3fn) - a["gemm2_weights"] = eye.expand(num_experts, width, width).contiguous() - a["gemm2_weights_scale"] = torch.ones( - num_experts, width // BLOCK, width // BLOCK, device=DEV, dtype=torch.float32 - ) - a["topk_weights"] = torch.ones(tokens, 1, device=DEV, dtype=torch.bfloat16) - out = op(**a) - exact = (op.reference(**a) == out).float().mean().item() - print(f" T={tokens} width={width}: recipe reproduces {100 * exact:.2f}%") - assert exact >= 0.999, f"modelled recipe reproduced only {100 * exact:.2f}%" - - assert torch.equal(_variant(a), op.reference(**a)), "the variant model drifted from reference" - wrong = { - "power-of-two scale": _variant(a, requant=_pow2_requantize), - "no activation requant": _variant(a, quant_act=False), - "no FC1-output requant": _variant(a, quant_fc1_out=False), - "gate/up swapped": _variant(a, swap=True), - } - for label, pred in wrong.items(): - frac = (pred == out).float().mean().item() - print(f" {label}: {100 * frac:.2f}%") - assert frac <= 0.10, f"{label} reproduced {100 * frac:.2f}% -- not discriminated" - - -# ── what the gate separates ─────────────────────────────────────────────── - - -def test_wrong_variants_land_outside_the_gate(): - """Every operand-level mistake, measured against one call at the checkpoint.""" - a = _args(1024, **CKPT, seed=77) - out = op(**a) - _compare(out, op.reference(**a)) - good = _ulp(out, op.reference(**a))[1] - for label, kw in ( - ("gate/up swapped", dict(swap=True)), - ("no FC1-output requant", dict(quant_fc1_out=False)), - ("no activation requant", dict(quant_act=False)), - ("neither requant", dict(quant_fc1_out=False, quant_act=False)), - ): - elt, rms = _ulp(out, _variant(a, **kw)) - print(f" {label}: {elt:.2f} ulp element-wise, {rms:.3f} ulp RMS") - # The RMS gate separates all four; the element gate does not separate - # a missing activation requant, hence the bit-exact pin above. - assert rms > 0.75, f"{label} is inside the RMS gate" - assert rms >= 2.0 * good, f"{label} is within 2x of the correct reference" - - -def test_transposed_activation_scale_is_a_silent_wrong_answer(): - """At `T == H/128` both of the op's dimension checks pass on the transpose.""" - a = _args(20, **CKPT, seed=99) - assert a["hidden_states_scale"].shape == (20, 20) - _compare(op(**a), op.reference(**a)) - swapped = op(**dict(a, hidden_states_scale=a["hidden_states_scale"].t().contiguous())) - elt, rms = _ulp(swapped, op.reference(**a)) - print(f" transposed: {elt:.2f} ulp element-wise, {rms:.3f} ulp RMS") - assert rms > 10.0, "the transposed layout must be visibly wrong" - - -# ── expert parallelism and the window ───────────────────────────────────── - - -def test_expert_parallel_windows_sum_to_the_full_stack(): - """Four 128-wide windows of 512: each matches its own reference, a token - with no slot in the window is zero, and the four sum to the full call.""" - full = _args(64, **CKPT, seed=7) - out = op(**full) - _compare(out, op.reference(**full)) - ids = full["topk_ids"].long() - summed = torch.zeros(out.shape, device=DEV) - for offset in (0, 128, 256, 384): - a = _args(64, **CKPT, local_expert_offset=offset, local_num_experts=128, seed=7) - assert torch.equal(a["topk_ids"], full["topk_ids"]) - got = op(**a) - _compare(got, op.reference(**a)) - dead = (got == 0).all(dim=-1) - expected_dead = ~((ids >= offset) & (ids < offset + 128)).any(dim=-1) - assert torch.equal(dead, expected_dead), "a token with no slot in the window must be zero" - summed += got.float() - # Each window rounds its partial to bf16 first: inside the gate, not bitwise. - _compare(summed.to(torch.bfloat16), op.reference(**full)) - - -def test_window_past_the_routing_space_is_all_zero(): - a = _args(64, 256, 128, 8, 2, local_expert_offset=8, local_num_experts=4, seed=22) - assert bool((op(**a) == 0).all()) - - -def test_repeated_expert_ids_are_counted_once_or_twice_by_token_count(): - """Counted once or twice depending on which routing kernel the token count - selects. Records the verdict per token count; asserts only that each call - matches one of the two readings, not that both readings occur.""" - seen = set() - for tokens in (16, 256): - a = _args(tokens, 256, 128, 8, 2, seed=tokens * 3 + 256) - a["topk_ids"][:, 1] = a["topk_ids"][:, 0] - out = op(**a) - _, twice = _ulp(out, _variant(a, count_repeats=True)) - _, once = _ulp(out, _variant(a, count_repeats=False)) - verdict = "once" if once < twice else "twice" - print(f" T={tokens}: counted {verdict} (once={once:.3f} twice={twice:.3f} ulp RMS)") - assert min(once, twice) <= 0.75, "neither reading of a repeat matches" - seen.add(verdict) - print(f" verdicts across token counts: {sorted(seen)}") - - -# ── arguments ───────────────────────────────────────────────────────────── - - -def test_gemm1_clamp_limit_is_live(): - a = _args(16, **CKPT, gemm1_clamp_limit=1.0, seed=16) - clamped = op(**a) - _compare(clamped, op.reference(**a)) - _, rms = _ulp(clamped, op.reference(**dict(a, gemm1_clamp_limit=None))) - print(f" clamp=1.0 vs the unclamped reference: {rms:.3f} ulp RMS") - assert rms > 0.75, "the clamp must change the result" - b = dict(a) - del b["gemm1_clamp_limit"] - assert torch.equal(op(**dict(a, gemm1_clamp_limit=None)), op(**b)) - - -def test_inert_arguments_are_bitwise_identical(): - a = _args(16, **CKPT, seed=16) - base = op(**a) - for label, kw in ( - ("routed_scaling_factor=1.0", dict(routed_scaling_factor=1.0)), - ("routed_scaling_factor=2.5", dict(routed_scaling_factor=2.5)), - ("routing_method_type=0", dict(routing_method_type=0)), - ("routing_method_type=2", dict(routing_method_type=2)), - ("routing_method_type=4", dict(routing_method_type=4)), - ("n_group=1, topk_group=1", dict(n_group=1, topk_group=1)), - ("tune_max_num_tokens=128", dict(tune_max_num_tokens=128)), - ("use_dp=True", dict(use_dp=True)), - ): - assert torch.equal(op(**dict(a, **kw)), base), f"{label} changed the result" +# ── what a cell cannot carry ────────────────────────────────────────────── def test_autotuned_tactic_matches_the_fallback_tactic(): @@ -398,68 +177,6 @@ def test_autotuned_tactic_matches_the_fallback_tactic(): AutoTuner.get().clear_cache() -def test_registered_fake_returns_the_wrong_hidden_size(): - """An upstream defect, recorded so nothing relies on it: the fake assumes - packed fp4 hidden states and doubles the width.""" - a = _args(16, 256, 128, 8, 2, seed=3) - real = op(**a) - with torch._subclasses.FakeTensorMode(): - fake_args = { - k: (torch.empty(v.shape, dtype=v.dtype, device=DEV) if torch.is_tensor(v) else v) - for k, v in a.items() - } - fake_shape = tuple(op(**fake_args).shape) - print(f" fake says {fake_shape}, the kernel returns {tuple(real.shape)}") - assert fake_shape == (real.shape[0], 2 * real.shape[1]) - - -def test_determinism_and_purity(): - a = _args(16, 512, 256, 16, 4, seed=25) - snaps = {k: v.clone() for k, v in a.items() if torch.is_tensor(v)} - assert torch.equal(op(**a), op(**a)) - for k, v in snaps.items(): - assert torch.equal(a[k], v), f"{k} was mutated" - - -def test_output_buffer_semantics(): - """`output=` is written in place and the call returns an empty `[0]` tensor; - rows of a taller buffer past `T` are untouched.""" - a = _args(16, 256, 128, 8, 2, seed=3) - fresh = op(**a) - buf = torch.full((32, 256), 7.0, device=DEV, dtype=torch.bfloat16) - ret = op(**dict(a, output=buf[:16])) - assert ret.shape == (0,) and ret.dtype == torch.bfloat16 - assert torch.equal(buf[:16], fresh) - assert bool((buf[16:] == 7.0).all()), "rows past T must be untouched" - - -def test_rejections(): - """The op's own loud checks; none of these needs a guard.""" - a = _args(16, 256, 128, 8, 2, seed=3) - - def refuses(label, **kw): - with pytest.raises(RuntimeError) as excinfo: - op(**dict(a, **kw)) - torch.cuda.synchronize() - print(f" {label}: {str(excinfo.value).splitlines()[0][:100]}") - - refuses("bf16 hidden_states", hidden_states=a["hidden_states"].to(torch.bfloat16)) - refuses("bf16 hidden_states_scale", hidden_states_scale=a["hidden_states_scale"].bfloat16()) - refuses("1-D hidden_states_scale", hidden_states_scale=a["hidden_states_scale"].flatten()) - refuses("fp16 weight scale", gemm1_weights_scale=a["gemm1_weights_scale"].half()) - refuses("int64 topk_ids", topk_ids=a["topk_ids"].long()) - refuses("fp32 topk_weights", topk_weights=a["topk_weights"].float()) - refuses("no routing given", topk_ids=None, topk_weights=None) - refuses("intermediate_size mismatch", intermediate_size=256) - refuses("num_experts == top_k", num_experts=2) - refuses("fused shared experts pre-routed", num_fused_shared_experts=1) - refuses("output of the wrong dtype", output=torch.empty(16, 256, device=DEV).half()) - refuses( - "output of the wrong shape", - output=torch.empty(16, 128, device=DEV, dtype=torch.bfloat16), - ) - - # ── the guards ──────────────────────────────────────────────────────────── @@ -471,39 +188,8 @@ def _strided(t: torch.Tensor) -> torch.Tensor: return view -def test_unguarded_hazards_are_silent(): - """What `is_valid` exists for: the raw op takes each of these and answers.""" - a = _args(16, 256, 128, 8, 2, seed=11) - good = op.reference(**a) - _, rms = _ulp(op(**dict(a, hidden_states=_strided(a["hidden_states"]))), good) - print(f" strided hidden_states accepted: {rms:.3f} ulp RMS") - assert rms > 10.0, "a strided operand must be visibly wrong, or the guard is pointless" - base = op(**a) - assert torch.equal(op(**dict(a, act_type=1)), base), "act_type is no longer ignored" - bias = torch.tensor([1e3, -1e3] * 4, device=DEV, dtype=torch.bfloat16) - assert torch.equal(op(**dict(a, routing_bias=bias)), base), ( - "routing_bias is no longer ignored on the pre-routed path" - ) - - -def test_guards_refuse(): +def test_guard_refuses_a_strided_hidden_states(): + """What `is_valid` exists for: the raw op reads a strided operand as dense.""" a = _args(16, 256, 128, 8, 2, seed=11) - with validating(op): - ok = op(**a) - for name in ( - "hidden_states", - "hidden_states_scale", - "gemm1_weights", - "gemm1_weights_scale", - "gemm2_weights", - "gemm2_weights_scale", - "topk_weights", - "topk_ids", - ): - with pytest.raises(AssertionError, match="must be contiguous"): - op(**dict(a, **{name: _strided(a[name])})) - with pytest.raises(AssertionError, match="silently ignored"): - op(**dict(a, act_type=1)) - with pytest.raises(AssertionError, match="routing_bias is ignored"): - op(**dict(a, routing_bias=torch.zeros(8, device=DEV, dtype=torch.bfloat16))) - assert torch.equal(op(**a), ok) + with validating(op), pytest.raises(AssertionError, match="must be contiguous"): + op(**dict(a, hidden_states=_strided(a["hidden_states"]))) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fused_sigmoid_gate_mul_add.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fused_sigmoid_gate_mul_add.py index 9eed4e2ec169..035ca9c348ce 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fused_sigmoid_gate_mul_add.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fused_sigmoid_gate_mul_add.py @@ -3,9 +3,8 @@ """GPU test for the fused_sigmoid_gate_mul_add catalog entry. The entry owns the contract. This file turns a cell into operands, locks both -sides of the in-place write, and holds the claims no cell carries: gate layouts -bitwise interchangeable, aliasing, the empty batch, the wrong-variant -discrimination and the loud rejections. +sides of the in-place write, and holds the claims no cell carries: the empty +batch and the loud rejection of a wrong-length gate. """ import zlib @@ -24,38 +23,27 @@ DEV = "cuda" HIDDEN = 2560 -PAD = 64 def _operands(spec, seed): - """`(final, gate, shared, output, buffers)`; `buffers` holds the padded - allocations a `dest="padded"` cell writes a window of.""" + """`(final, gate, shared)`: the in-place combine's operands.""" t, h, dt = spec["tokens"], spec["hidden"], spec["dtype"] + assert spec["dest"] == "inplace", spec["dest"] g = torch.Generator(device=DEV).manual_seed(seed) def rand(*shape, dtype=dt): return torch.randn(*shape, device=DEV, generator=g).to(dtype) - buffers = {} - if spec["dest"] == "padded": - buffers = {k: rand(t + 2, h + PAD) for k in ("final", "shared", "output")} - final, shared, output = (buffers[k][:t, :h] for k in ("final", "shared", "output")) - assert output.stride() == (h + PAD, 1) and not output.is_contiguous() - else: - final, shared = rand(t, h), rand(t, h) - output = torch.empty_like(final) if spec["dest"] == "output" else None + final, shared = rand(t, h), rand(t, h) gate = rand(t, 1, dtype=spec["gate_dtype"]) - layout = spec["gate_layout"] - if layout == "flat": - gate = gate.reshape(-1) - elif layout == "row": - gate = gate.reshape(1, t) - elif layout == "strided": + if spec["gate_layout"] == "strided": wide = rand(t, 4, dtype=spec["gate_dtype"]) wide[:, 2:3] = gate gate = wide[:, 2:3] assert not gate.is_contiguous() - return final, gate, shared, output, buffers + else: + assert spec["gate_layout"] == "column", spec["gate_layout"] + return final, gate, shared # ── the certified cells ─────────────────────────────────────────────────── @@ -64,68 +52,21 @@ def rand(*shape, dtype=dt): @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name certifying(op) - final, gate, shared, output, buffers = _operands(cell.spec, zlib.crc32(cell.why.encode())) + final, gate, shared = _operands(cell.spec, zlib.crc32(cell.why.encode())) before = {k: v.clone() for k, v in dict(final=final, gate=gate, shared=shared).items()} - snapshots = {k: v.clone() for k, v in buffers.items()} - ref = op.reference(before["final"], gate, shared, output) + ref = op.reference(before["final"], gate, shared) - out = op(final, gate, shared, output) + out = op(final, gate, shared) - destination = final if output is None else output - assert out.data_ptr() == destination.data_ptr(), "the op returns what it wrote" + assert out.data_ptr() == final.data_ptr(), "the op returns what it wrote" assert out.dtype == final.dtype and out.shape == final.shape torch.testing.assert_close(out, ref) assert torch.equal(gate, before["gate"]) and torch.equal(shared, before["shared"]) - if output is not None: - assert torch.equal(final, before["final"]), "an output buffer leaves the input alone" - if buffers: - t, h = cell.spec["tokens"], cell.spec["hidden"] - o, o0 = buffers["output"], snapshots["output"] - assert torch.equal(o[:t, h:], o0[:t, h:]), "wrote into the row padding" - assert torch.equal(o[t:], o0[t:]), "wrote past num_tokens" - assert torch.equal(buffers["final"], snapshots["final"]) - assert torch.equal(buffers["shared"], snapshots["shared"]) # ── what a cell cannot carry ────────────────────────────────────────────── -def test_gate_layouts_are_the_same_call_bit_for_bit(): - spec = dict( - tokens=33, - hidden=HIDDEN, - dtype=torch.bfloat16, - gate_dtype=torch.bfloat16, - gate_layout="column", - dest="output", - ) - final, gate, shared, _, _ = _operands(spec, 3) - wide = torch.zeros(33, 4, dtype=torch.bfloat16, device=DEV) - wide[:, 2:3] = gate - results = [ - op(final.clone(), g, shared) - for g in (gate, gate.reshape(-1), gate.reshape(1, 33), wide[:, 2:3]) - ] - for r in results[1:]: - assert torch.equal(r, results[0]) - - -def test_the_output_buffer_may_alias_the_shared_expert_output(): - spec = dict( - tokens=17, - hidden=HIDDEN, - dtype=torch.bfloat16, - gate_dtype=torch.bfloat16, - gate_layout="column", - dest="inplace", - ) - final, gate, shared, _, _ = _operands(spec, 8) - expected = op.reference(final, gate, shared) - returned = op(final, gate, shared, output=shared) - assert returned.data_ptr() == shared.data_ptr() - torch.testing.assert_close(shared, expected) - - def test_zero_tokens_is_a_no_op(): final = torch.empty(0, HIDDEN, dtype=torch.bfloat16, device=DEV) shared = torch.empty(0, HIDDEN, dtype=torch.bfloat16, device=DEV) @@ -134,67 +75,9 @@ def test_zero_tokens_is_a_no_op(): assert returned.shape == (0, HIDDEN) and returned.data_ptr() == final.data_ptr() -def test_wrong_variants_land_far_outside_the_default_gate(): - """The default bf16 band (atol 1e-5, rtol 1.6e-2) separates the correct - combine from every plausible wrong one by more than 100x.""" - spec = dict( - tokens=17, - hidden=HIDDEN, - dtype=torch.bfloat16, - gate_dtype=torch.bfloat16, - gate_layout="column", - dest="inplace", - ) - final, gate, shared, _, _ = _operands(spec, 9) - correct = op.reference(final, gate, shared) - actual = op(final.clone(), gate, shared) - width = 1e-5 + 1.6e-2 * correct.float().abs() - s = torch.sigmoid(gate.reshape(-1, 1).float()) - f, sh = final.float(), shared.float() - - def ratio(wrong: torch.Tensor) -> float: - return ((wrong.float() - correct.float()).abs() / width).max().item() - - assert ratio(actual) <= 1.0 - for name, wrong in { - "no sigmoid, raw gate": f + gate.float().reshape(-1, 1) * sh, - "gate applied to the routed half": s * f + sh, - "one gate for the whole batch": f + s.reshape(-1)[0] * sh, - "multiply instead of add": f * s * sh, - "gate dropped, plain add": f + sh, - }.items(): - r = ratio(wrong.to(torch.bfloat16)) - print(f" {name:36s} {r:10.1f}x the gate") - assert r > 100.0, f"{name} is indistinguishable from the correct result" - - -@pytest.mark.parametrize( - "case", - ["shared_shape_mismatch", "gate_numel_too_small", "final_last_stride_without_an_output"], -) -def test_metadata_misuse_raises(case: str) -> None: - """Metadata violations are loud, which is why the entry has no guard: one - case each for an operand that disagrees with `final_hidden_states`, a gate - of the wrong length and a destination whose last stride is not 1.""" +def test_a_gate_of_the_wrong_length_raises() -> None: + """Metadata violations are loud, which is why the entry has no guard.""" bf = dict(dtype=torch.bfloat16, device=DEV) - final, shared, gate = ( - torch.randn(4, 128, **bf), - torch.randn(4, 128, **bf), - torch.randn(4, 1, **bf), - ) - if case == "shared_shape_mismatch": - shared = torch.randn(4, 64, **bf) - elif case == "gate_numel_too_small": - gate = torch.randn(3, 1, **bf) - elif case == "final_last_stride_without_an_output": - final = torch.randn(4, 256, **bf)[:, ::2] + final, shared = torch.randn(4, 128, **bf), torch.randn(4, 128, **bf) with pytest.raises(AssertionError): - op(final, gate, shared) - - -def test_cpu_operands_are_rejected_by_triton() -> None: - final = torch.randn(4, 128, dtype=torch.bfloat16) - shared = torch.randn(4, 128, dtype=torch.bfloat16) - gate = torch.randn(4, 1, dtype=torch.bfloat16) - with pytest.raises(ValueError, match="Pointer argument cannot be accessed from Triton"): - op(final, gate, shared) + op(final, torch.randn(3, 1, **bf), shared) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_sort.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_sort.py index 266a8a23a9d1..8d97d4b0dd4f 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_sort.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_sort.py @@ -3,10 +3,8 @@ """GPU test for the moe_sort entry of the Qwen3.8-Flash-Next target. The entry owns the contract. This file turns a cell into expert ids and holds -the claims no cell carries: CUDA-graph replay at the captured row counts, the -eager row sweep, which regions the op leaves unwritten, the token-order layout -of the block kernel, the scales being inert, the data hazards, the loud -rejections and the guards. +the claims no cell carries: CUDA-graph replay at the largest captured row +count, and the contiguity guard. """ import zlib @@ -25,7 +23,6 @@ DEV = "cuda" E, K = 512, 10 -SENTINEL = 0x5A5A5A5A _NAMES = ( "tile_idx_to_expert_idx", @@ -91,10 +88,6 @@ def _ids(dist, tokens, tile, seed, num_experts=E, top_k=K): if dist == "router": scores = torch.rand(tokens, num_experts, device=DEV, generator=g) return scores.topk(top_k, dim=-1).indices.to(torch.int32) - if dist == "hot": - hot = torch.randperm(num_experts, device=DEV, generator=g)[:top_k] - slots = torch.rand(tokens, top_k, device=DEV, generator=g).argsort(dim=-1) - return hot[slots].to(torch.int32) if dist == "tile_edges": edge = (7, 300, 511) scores = torch.rand(tokens, num_experts, device=DEV, generator=g) @@ -103,12 +96,6 @@ def _ids(dist, tokens, tile, seed, num_experts=E, top_k=K): for slot, (expert, rows) in enumerate(zip(edge, (tile - 1, tile, tile + 1))): ids[:rows, slot] = expert return ids - if dist == "round_robin": - base = torch.arange(tokens, device=DEV).unsqueeze(1) * top_k - return ((base + torch.arange(top_k, device=DEV)) % num_experts).to(torch.int32) - if dist == "strided": - h = torch.randint(0, num_experts, (tokens, 1), device=DEV, generator=g) - return ((h + torch.arange(top_k, device=DEV) * 51) % num_experts).to(torch.int32) raise ValueError(dist) @@ -157,9 +144,6 @@ def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves th out, ref = _check(args) assert torch.equal(ids, ids_before), "token_selected_experts was written" assert torch.equal(args["token_final_scales"], scales_before), "token_final_scales was written" - if s["dist"] == "round_robin": - # every pair on its own expert: the shape bound is reached exactly - assert int(out[5]) == out[0].numel() == s["tokens"] * s["top_k"] if s["dist"] == "tile_edges": tile = s["tile_tokens_dim"] n = int(ref[5]) @@ -176,10 +160,8 @@ def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves th # ── CUDA graphs: captured row counts, on new ids ────────────────────────── -#: Row counts of a decode capture (requests x MTP verify width 1..4, up to 32 -#: requests): both ends of the block kernel (T <= 4), the first cluster-kernel T, -#: and the largest capture. -_CAPTURED_ROWS = (1, 4, 5, 128) +#: The largest decode capture: 32 requests at MTP verify width 4, cluster kernel. +_CAPTURED_ROWS = 128 def _capture(args): @@ -197,234 +179,37 @@ def _capture(args): def test_captured_rows_replay_on_new_ids(certifying): # noqa: F811 - """Each graph is replayed on two fresh id sets and gated against the reference, - and against an eager call on the same ids: counts and tile arrays bitwise, the - whole output bitwise where the kernel is deterministic (T <= 4). At the - target's tile, 128.""" + """The graph is replayed on two fresh id sets and gated against the reference, + and against an eager call on the same ids: counts and tile arrays bitwise. At + the target's tile, 128.""" certifying(op) - tile = 128 - for rows in _CAPTURED_ROWS: - args = _args(_ids("router", rows, tile, seed=rows), tile, seed=rows) - graph, static_out = _capture(args) - for trial in (1, 2): - fresh = _ids("router", rows, tile, seed=1000 * trial + rows) - args["token_selected_experts"].copy_(fresh) - graph.replay() - eager = op(**dict(args, token_selected_experts=fresh.clone())) - ref = op.reference(**args) - _compare(static_out, ref) - _compare(eager, ref) - n = int(ref[5]) - for i in (0, 1): - assert torch.equal(static_out[i][:n], eager[i][:n]), (rows, trial, i) - for i in (4, 5): - assert torch.equal(static_out[i], eager[i]), (rows, trial, i) - if rows <= 4: - real = ref[3] >= 0 - assert torch.equal(static_out[2], eager[2]), (rows, trial) - assert torch.equal(static_out[3][real], eager[3][real]), (rows, trial) - del graph, static_out - torch.cuda.synchronize() - - -@pytest.mark.parametrize("rows", (3, 300)) -def test_padding_rows_and_spare_tiles_are_not_written(rows): - """A sentinel planted in a captured graph's outputs survives replay in exactly - the regions `_compare` excludes; under a graph those regions keep the previous - replay's indices. Block kernel (3 rows) and cluster kernel (300), tile 128.""" - tile = 128 - args = _args(_ids("router", rows, tile, seed=5), tile, seed=5) + tile, rows = 128, _CAPTURED_ROWS + args = _args(_ids("router", rows, tile, seed=rows), tile, seed=rows) graph, static_out = _capture(args) for trial in (1, 2): - args["token_selected_experts"].copy_(_ids("router", rows, tile, seed=50 + trial)) - for t in static_out[:4]: - t.fill_(SENTINEL) + fresh = _ids("router", rows, tile, seed=1000 * trial + rows) + args["token_selected_experts"].copy_(fresh) graph.replay() + eager = op(**dict(args, token_selected_experts=fresh.clone())) ref = op.reference(**args) _compare(static_out, ref) + _compare(eager, ref) n = int(ref[5]) - real = ref[3] >= 0 - assert bool((static_out[0][n:] == SENTINEL).all()) and bool( - (static_out[1][n:] == SENTINEL).all() - ) - assert bool((static_out[3][~real] == SENTINEL).all()) - assert not bool((static_out[3][real] == SENTINEL).any()) - assert not bool((static_out[2] == SENTINEL).any()), ( - "expanded_idx_to_permuted_idx is fully written" - ) - # replay on other ids without re-planting: the unwritten rows keep the last replay's values - before = static_out[3].clone() - args["token_selected_experts"].copy_(_ids("router", rows, tile, seed=99)) - graph.replay() - real = op.reference(**args)[3] >= 0 - assert torch.equal(static_out[3][~real], before[~real]) - if rows == 300: - stale = before[~real] - assert int(((stale >= 0) & (stale < rows * K)).sum()) > 0, "no valid-looking stale row" + for i in (0, 1): + assert torch.equal(static_out[i][:n], eager[i][:n]), (trial, i) + for i in (4, 5): + assert torch.equal(static_out[i], eager[i]), (trial, i) del graph, static_out + torch.cuda.synchronize() -# ── eager: the row sweep ────────────────────────────────────────────────── - -#: Both sides of every kernel-path edge: block / cluster at 4 / 5, the cluster -#: kernel's per-block work edges at 102 / 103 and 204 / 205 rows, its fast / slow -#: path edges at 3276 / 3277 and 6553 / 6554, and the 8192-row ceiling. -_SWEEP = (4, 5, 102, 103, 204, 205, 3276, 3277, 6553, 6554, 8192) +# ── the guard ───────────────────────────────────────────────────────────── -def test_eager_row_sweep(certifying): # noqa: F811 - certifying(op) - tile = 128 - for tokens in _SWEEP: - _check(_args(_ids("router", tokens, tile, seed=tokens), tile, seed=tokens)) - - -# ── layout facts ────────────────────────────────────────────────────────── - - -def test_block_kernel_rows_follow_token_order(): - """T <= 4: each expert's rows in token order, bitwise the reference's.""" - for tile in (128, 256): - for tokens in (1, 4): - for seed in range(2): - args = _args(_ids("router", tokens, tile, seed=seed), tile, seed=seed) - out, ref = _check(args) - real = ref[3] >= 0 - assert torch.equal(out[2], ref[2]) and torch.equal(out[3][real], ref[3][real]) - - -def test_scales_are_inert_and_untouched(): - """The scales' values and dtype never reach an output, and are not written.""" - for tokens in (3, 300): - ids = _ids("router", tokens, 128, seed=21) - outs = [] - for scales in ( - torch.rand(tokens, K, device=DEV), - torch.rand(tokens, K, device=DEV).to(torch.bfloat16), - torch.zeros(tokens, K, device=DEV), - torch.full((tokens, K), float("nan"), device=DEV), - ): - bits = scales.view( - torch.int32 if scales.dtype == torch.float32 else torch.int16 - ).clone() - args = dict(_args(ids, 128), token_final_scales=scales) - out, ref = _check(args) - assert torch.equal(scales.view(bits.dtype), bits), "token_final_scales was written" - outs.append(out) - if tokens <= 4: # deterministic: every written element bitwise across the four - n, real = int(ref[5]), ref[3] >= 0 - for out in outs[1:]: - assert torch.equal(out[0][:n], outs[0][0][:n]) and torch.equal( - out[1][:n], outs[0][1][:n] - ) - assert torch.equal(out[2], outs[0][2]) and torch.equal( - out[3][real], outs[0][3][real] - ) - assert torch.equal(out[4], outs[0][4]) and torch.equal(out[5], outs[0][5]) - - -# ── data hazards is_valid cannot see ────────────────────────────────────── - - -def test_duplicate_expert_in_a_row(): - """Block kernel (T <= 4) keeps one of two slots naming the same expert; - the other paths keep both.""" - ids = _ids("router", 2, 128, seed=31) - ids[:, 1] = ids[:, 0] - args = _args(ids, 128) - out = op(**args) - ref = op.reference(**args) - for row in range(2): - pair = out[2][row, :2].tolist() - assert pair.count(-1) == 1, pair - assert int((out[2] >= 0).sum()) == 2 * K - 2 and int((ref[2] >= 0).sum()) == 2 * K - with pytest.raises(AssertionError): - _compare(out, ref) - - ids = _ids("router", 300, 128, seed=32) - ids[:, 1] = ids[:, 0] - _check(_args(ids, 128)) - - -def test_ids_past_int16_alias_off_the_block_kernel(): - """An id of 65536 + e is not local, and the block kernel says so; the - cluster kernel reads ids through int16 and sorts it as expert e.""" - for tokens in (2, 300): - ids = _ids("router", tokens, 128, seed=33) - ids[0, 0] = 40000 # negative as int16: non-local everywhere - out, _ = _check(_args(ids, 128)) - assert int(out[2][0, 0]) == -1 - ids[0, 0] = 65536 + 3 - out = op(**_args(ids, 128)) - p = int(out[2][0, 0]) - if tokens <= 4: - assert p == -1 - _compare(out, op.reference(**_args(ids, 128))) - else: - assert p >= 0 and int(out[0][p // 128]) == 3 - with pytest.raises(AssertionError): - _compare(out, op.reference(**_args(ids, 128))) - - -# ── loud rejections ─────────────────────────────────────────────────────── - - -def test_op_rejects_loudly(): - """Rejected by the op itself; none of these needs a guard.""" - ids = _ids("router", 64, 128, seed=41) - sc = torch.rand(64, K, device=DEV) - wide33 = _ids("router", 8, 128, seed=42, top_k=33) - cases = [ - ("int64 ids", lambda: op(ids.long(), sc, E, K, 0, E, 128), "expected scalar type Int"), - ("1-D ids", lambda: op(ids.reshape(-1), sc, E, K, 0, E, 128), "must be 2D"), - ("top_k vs width", lambda: op(ids, sc, E, 9, 0, E, 128), r"size\(1\) must be top_k"), - ("scales width", lambda: op(ids, sc[:, :9], E, K, 0, E, 128), r"size\(1\) must be top_k"), - ("scales rows", lambda: op(ids, sc[:63], E, K, 0, E, 128), "must be num_tokens"), - ("top_k 33", lambda: op(wide33, torch.rand(8, 33, device=DEV), E, 33, 0, E, 128), "<= 32"), - ("num_experts 2049", lambda: op(ids, sc, 2049, K, 0, 2049, 128), "no more than 2048"), - ("T = 0", lambda: op(ids[:0], sc[:0], E, K, 0, E, 128), "mPtrTopKWeights must also"), - ] - for tag, call, match in cases: - with pytest.raises(RuntimeError, match=match): - call() - with pytest.raises(NotImplementedError): - op(ids.cpu(), sc.cpu(), E, K, 0, E, 128) - - -# ── the guards ──────────────────────────────────────────────────────────── - - -def test_unguarded_noncontiguous_ids_answer_wrongly(): +def test_guard_refuses_noncontiguous_ids(): """What `is_valid` exists for: the op reads a strided view as the dense buffer that starts at its data pointer, and answers without raising.""" wide = _ids("router", 300, 128, seed=51, top_k=16) - view = wide[:, :K] - sc = torch.rand(300, K, device=DEV) - out = op(view, sc, E, K, 0, E, 128) - dense = wide.reshape(-1)[: 300 * K].reshape(300, K) - _compare(out, op.reference(dense, sc, E, K, 0, E, 128)) - with pytest.raises(AssertionError): - _compare(out, op.reference(view.contiguous(), sc, E, K, 0, E, 128)) - - -def test_guards_refuse(): - """The 2^31-element operands are expanded views of one row: the guard reads - only their shape, so nothing is allocated.""" - wide = _ids("router", 300, 128, seed=51, top_k=16) sc = torch.rand(300, K, device=DEV) - ids = _ids("router", 300, 128, seed=53) - base = ids[:1] - with validating(op): - with pytest.raises(AssertionError, match="contiguous"): - op(wide[:, :K], sc, E, K, 0, E, 128) - with pytest.raises(AssertionError, match="contiguous"): - op(ids.t().contiguous().t(), sc, E, K, 0, E, 128) - for tokens in (214748365, 429496730): - with pytest.raises(AssertionError, match=r"2\^31"): - op(base.expand(tokens, K), sc[:1].expand(tokens, K), E, K, 0, E, 128) - # tile_tokens_dim = 0 is refused but not driven unguarded: the host divides - # by it, which returns garbage on aarch64 and traps on x86. - with pytest.raises(AssertionError, match="tile_tokens_dim must be positive"): - op(ids, sc, E, K, 0, E, 0) - out = op(wide[:, :K].contiguous(), sc, E, K, 0, E, 128) - _compare(out, op.reference(wide[:, :K].contiguous(), sc, E, K, 0, E, 128)) + with validating(op), pytest.raises(AssertionError, match="contiguous"): + op(wide[:, :K], sc, E, K, 0, E, 128) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_unpermute_inplace.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_unpermute_inplace.py index ef20128ffb78..1ff8171c9b9e 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_unpermute_inplace.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_unpermute_inplace.py @@ -5,9 +5,8 @@ The entry owns the contract. This file lays a cell's FC2 rows out the way the CuteDSL FC2 hands them over -- placed by moe_sort, every row the op must not read filled with NaN, the output pre-filled with NaN -- and holds the claims no cell -carries: CUDA-graph replay at the captured row counts, the eager prefix sweep, -run-to-run and placement independence, the gate's discrimination, T = 0, the loud -rejections and the guards. +carries: CUDA-graph replay at the largest captured row count, bitwise equal to +eager, and the contiguity guard. """ import zlib @@ -28,13 +27,8 @@ E, K, H = 512, 10, 2560 NAN = float("nan") -#: Row counts of a decode capture (requests x MTP verify width 1..4, up to 32 -#: requests): both ends of the sort's block kernel, its first cluster-kernel T and -#: the largest capture. The eager prefix lengths straddle the captured ceiling and -#: reach the 8192-row one; a subset of them is also re-sorted. -_CAPTURED_ROWS = (1, 4, 5, 128) -_PREFIXES = (1, 4, 5, 127, 128, 129, 4096, 8191) -_RESORT = (1, 5, 129, 8192) +#: The largest decode capture: 32 requests at MTP verify width 4. +_CAPTURED_ROWS = 128 def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: @@ -62,31 +56,16 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: ) -def _spec(tokens, tile_tokens_dim=128, dist="router", values="normal", non_local_from=None): +def _spec(tokens, tile_tokens_dim=128, values="normal"): """The fields `_problem` reads; the same names as a cell's spec.""" - return dict( - tokens=tokens, - tile_tokens_dim=tile_tokens_dim, - dist=dist, - values=values, - non_local_from=non_local_from, - ) + return dict(tokens=tokens, tile_tokens_dim=tile_tokens_dim, values=values) -def _ids(dist: str, tokens: int, seed: int) -> torch.Tensor: - """Expert ids, `[tokens, K]` int32, every row naming 10 distinct experts.""" +def _ids(tokens: int, seed: int) -> torch.Tensor: + """Expert ids, `[tokens, K]` int32: each row's top-10 of uniform scores.""" g = torch.Generator(device=DEV).manual_seed(seed) - if dist == "router": - scores = torch.rand(tokens, E, device=DEV, generator=g) - return scores.topk(K, dim=-1).indices.to(torch.int32) - if dist == "hot": - hot = torch.randperm(E, device=DEV, generator=g)[:K] - slots = torch.rand(tokens, K, device=DEV, generator=g).argsort(dim=-1) - return hot[slots].to(torch.int32) - if dist == "round_robin": - base = torch.arange(tokens, device=DEV).unsqueeze(1) * K - return ((base + torch.arange(K, device=DEV)) % E).to(torch.int32) - raise ValueError(dist) + scores = torch.rand(tokens, E, device=DEV, generator=g) + return scores.topk(K, dim=-1).indices.to(torch.int32) def _scales(tokens: int, seed: int) -> torch.Tensor: @@ -130,19 +109,16 @@ def _place(rows: torch.Tensor, e2p: torch.Tensor, num_rows: int) -> torch.Tensor return perm -def _output(tokens: int, dtype: torch.dtype = torch.bfloat16) -> torch.Tensor: - return torch.full((tokens, H), NAN, dtype=dtype, device=DEV) +def _output(tokens: int) -> torch.Tensor: + return torch.full((tokens, H), NAN, dtype=torch.bfloat16, device=DEV) def _problem(spec: dict, seed: int): """A cell's arguments, plus the per-pair rows and ids they were built from.""" tokens = spec["tokens"] - ids = _ids(spec["dist"], tokens, seed) + ids = _ids(tokens, seed) scales = _scales(tokens, seed) e2p, num_rows = _sort(ids, scales, spec["tile_tokens_dim"]) - if spec["non_local_from"] is not None: - e2p = torch.where(ids >= spec["non_local_from"], -1, e2p).to(torch.int32) - e2p[0] = -1 rows = _rows(tokens, seed, spec["values"]) args = dict( permuted_input=_place(rows, e2p, num_rows), @@ -150,7 +126,7 @@ def _problem(spec: dict, seed: int): expanded_idx_to_permuted_idx=e2p, topk_scales=scales, ) - return args, rows, ids + return args def _bits(t: torch.Tensor) -> torch.Tensor: @@ -163,8 +139,7 @@ def _same_bits(a: torch.Tensor, b: torch.Tensor) -> bool: def _cell_id(cell): s = cell.spec - tag = f"t{s['tile_tokens_dim']}-T{s['tokens']}-{s['dist']}-{s['values']}" - return tag if s["non_local_from"] is None else f"{tag}-nonlocal{s['non_local_from']}" + return f"t{s['tile_tokens_dim']}-T{s['tokens']}-{s['values']}" # ── the certified cells ─────────────────────────────────────────────────── @@ -176,7 +151,7 @@ def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves th s = cell.spec assert (s["num_experts"], s["top_k"], s["hidden_size"]) == (E, K, H) assert (s["dtype"], s["scales_dtype"]) == (torch.bfloat16, torch.float32) - args, _, _ = _problem(s, zlib.crc32(cell.why.encode())) + args = _problem(s, zlib.crc32(cell.why.encode())) before = { n: args[n].clone() for n in ("permuted_input", "expanded_idx_to_permuted_idx", "topk_scales") @@ -186,9 +161,6 @@ def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves th _compare(out, op.reference(**args)) for name, snapshot in before.items(): assert _same_bits(args[name], snapshot), f"{name} was written" - if s["non_local_from"] is not None: - assert int((args["expanded_idx_to_permuted_idx"] < 0).sum()) > s["tokens"] - assert bool((_bits(out[0]) == 0).all()), "a token with every slot -1 is not +0" # ── CUDA graphs: captured row counts, on new data ───────────────────────── @@ -210,7 +182,7 @@ def _capture(fn): def test_captured_rows_replay_on_new_data(certifying): # noqa: F811 """The sort, the FC2 rows' placement and this op are captured together, as the - target runs them; each graph is replayed on two fresh sets of ids, scales and + target runs them; the graph is replayed on two fresh sets of ids, scales and rows. The permuted buffer is not cleared between replays, so the rows the new sort does not use hold the previous placement's rows. Each replay is gated against the reference, and is bitwise equal to an eager call on the same data @@ -219,259 +191,55 @@ def test_captured_rows_replay_on_new_data(certifying): # noqa: F811 certifying(op) tile = 128 stale_rows = 0 - for tokens in _CAPTURED_ROWS: - ids = _ids("router", tokens, seed=tokens) - scales = _scales(tokens, seed=tokens) - rows = _rows(tokens, seed=tokens) - num_rows = _sort(ids, scales, tile)[1] - perm = torch.full((num_rows, H), NAN, dtype=torch.bfloat16, device=DEV) - out = _output(tokens) - - def step(ids=ids, scales=scales, rows=rows, perm=perm, out=out, tile=tile): - e2p = _sort(ids, scales, tile)[0] - perm.index_copy_(0, e2p.reshape(-1).long(), rows.reshape(-1, H)) - op( - permuted_input=perm, - output=out, - expanded_idx_to_permuted_idx=e2p, - topk_scales=scales, - ) - return e2p - - graph, e2p = _capture(step) - for trial in (1, 2): - seed = 1000 * trial + tokens - ids.copy_(_ids("router", tokens, seed)) - scales.copy_(_scales(tokens, seed)) - rows.copy_(_rows(tokens, seed)) - out.fill_(NAN) - graph.replay() - args = dict( - permuted_input=perm, - output=out, - expanded_idx_to_permuted_idx=e2p, - topk_scales=scales, - ) - _compare(out, op.reference(**args)) - used = torch.zeros(num_rows, dtype=torch.bool, device=DEV) - used[e2p.reshape(-1).long()] = True - stale_rows += int((~used & ~perm.isnan().any(1)).sum()) - - e2p_eager, n = _sort(ids.clone(), scales.clone(), tile) - eager = op( - permuted_input=_place(rows.clone(), e2p_eager, n), - output=_output(tokens), - expanded_idx_to_permuted_idx=e2p_eager, - topk_scales=scales.clone(), - ) - assert _same_bits(out, eager), (tokens, trial) - del graph, e2p - torch.cuda.synchronize() - assert stale_rows > 0, "no replay ran over a stale row: the stale-row claim is untested" - - -# ── eager: prefixes of the largest batch ────────────────────────────────── - - -def test_eager_rows_prefix_bitwise(certifying): # noqa: F811 - """One block per token makes the grid the only thing T sizes: the 8192-row call - is gated against the reference, and the first T rows of every row count in the - sweep are bitwise those of the 8192-row call -- over the same placement, and - (a subset) over a fresh sort of the first T tokens alone, which places their - rows elsewhere. A token's output does not depend on its batch. - """ - certifying(op) - tile = 128 - args, rows, ids = _problem(_spec(8192, tile_tokens_dim=tile), seed=8192 + tile) - full = op(**args) - _compare(full, op.reference(**args)) - perm, e2p, scales = ( - args["permuted_input"], - args["expanded_idx_to_permuted_idx"], - args["topk_scales"], - ) - for tokens in _PREFIXES: - out = op( + tokens = _CAPTURED_ROWS + ids = _ids(tokens, seed=tokens) + scales = _scales(tokens, seed=tokens) + rows = _rows(tokens, seed=tokens) + num_rows = _sort(ids, scales, tile)[1] + perm = torch.full((num_rows, H), NAN, dtype=torch.bfloat16, device=DEV) + out = _output(tokens) + + def step(ids=ids, scales=scales, rows=rows, perm=perm, out=out, tile=tile): + e2p = _sort(ids, scales, tile)[0] + perm.index_copy_(0, e2p.reshape(-1).long(), rows.reshape(-1, H)) + op( permuted_input=perm, - output=_output(tokens), - expanded_idx_to_permuted_idx=e2p[:tokens], - topk_scales=scales[:tokens], + output=out, + expanded_idx_to_permuted_idx=e2p, + topk_scales=scales, + ) + return e2p + + graph, e2p = _capture(step) + for trial in (1, 2): + seed = 1000 * trial + tokens + ids.copy_(_ids(tokens, seed)) + scales.copy_(_scales(tokens, seed)) + rows.copy_(_rows(tokens, seed)) + out.fill_(NAN) + graph.replay() + args = dict( + permuted_input=perm, + output=out, + expanded_idx_to_permuted_idx=e2p, + topk_scales=scales, ) - assert _same_bits(out, full[:tokens]), tokens - for tokens in _RESORT: - e2p_t, n = _sort(ids[:tokens], scales[:tokens], tile) - out = op( - permuted_input=_place(rows[:tokens], e2p_t, n), + _compare(out, op.reference(**args)) + used = torch.zeros(num_rows, dtype=torch.bool, device=DEV) + used[e2p.reshape(-1).long()] = True + stale_rows += int((~used & ~perm.isnan().any(1)).sum()) + + e2p_eager, n = _sort(ids.clone(), scales.clone(), tile) + eager = op( + permuted_input=_place(rows.clone(), e2p_eager, n), output=_output(tokens), - expanded_idx_to_permuted_idx=e2p_t, - topk_scales=scales[:tokens], + expanded_idx_to_permuted_idx=e2p_eager, + topk_scales=scales.clone(), ) - assert _same_bits(out, full[:tokens]), tokens - - -# ── run to run, and where the rows sit ──────────────────────────────────── - - -@pytest.mark.parametrize("tokens", (1, 8192)) -def test_repeats_and_placements_are_bitwise(tokens, certifying): # noqa: F811 - """What the fused finalize does not give: 8 identical calls, a second sort - of the same ids (the sort orders an expert's rows by atomics), a random - permutation of every permuted row, and a shuffle of the tokens all give the - same bits. - """ - certifying(op) - args, rows, ids = _problem(_spec(tokens), seed=tokens + 7) - first = op(**args).clone() - _compare(first, op.reference(**args)) - for _ in range(7): - assert _same_bits(op(**dict(args, output=_output(tokens))), first) - - perm, e2p, scales = ( - args["permuted_input"], - args["expanded_idx_to_permuted_idx"], - args["topk_scales"], - ) - e2p2, n = _sort(ids, scales, 128) - if tokens == 8192: - assert int((e2p2 != e2p).sum()) > 0, "the second sort placed every row where the first did" - out = op( - permuted_input=_place(rows, e2p2, n), - output=_output(tokens), - expanded_idx_to_permuted_idx=e2p2, - topk_scales=scales, - ) - assert _same_bits(out, first), "a second sort of the same ids changed the result" - - g = torch.Generator(device=DEV).manual_seed(tokens) - p = torch.randperm(perm.shape[0], device=DEV, generator=g) - moved = torch.empty_like(perm) - moved[p] = perm - out = op( - permuted_input=moved, - output=_output(tokens), - expanded_idx_to_permuted_idx=p[e2p.long()].to(torch.int32), - topk_scales=scales, - ) - assert _same_bits(out, first), "a random placement of the rows changed the result" - - order = torch.randperm(tokens, device=DEV, generator=g) - out = op( - permuted_input=perm, - output=_output(tokens), - expanded_idx_to_permuted_idx=e2p[order].contiguous(), - topk_scales=scales[order].contiguous(), - ) - assert _same_bits(out, first[order]), "shuffling the tokens changed a token's result" - - -# ── what the bitwise gate can see ───────────────────────────────────────── - - -def test_the_gate_sees_slot_order_and_rounding(certifying): # noqa: F811 - """`_compare` is bitwise; a sum taken in another order or with other roundings - is rejected. The reverse-order sum is the entry's own reference over the - slots reversed. - """ - certifying(op) - args, rows, _ = _problem(_spec(8192), seed=99) - out = op(**args) - _compare(out, op.reference(**args)) - scales = args["topk_scales"] - reverse = op.reference( - permuted_input=args["permuted_input"], - output=args["output"], - expanded_idx_to_permuted_idx=args["expanded_idx_to_permuted_idx"].flip(1).contiguous(), - topk_scales=scales.flip(1).contiguous(), - ) - unfused = torch.zeros(8192, H, device=DEV) - for k in range(K): - unfused = unfused + rows[:, k].float() * scales[:, k : k + 1] - exact = (rows.double() * scales.double().unsqueeze(-1)).sum(1) - bf16_sum = torch.zeros(8192, H, dtype=torch.bfloat16, device=DEV) - for k in range(K): - bf16_sum = bf16_sum + (rows[:, k].float() * scales[:, k : k + 1]).to(torch.bfloat16) - for wrong in ( - reverse, # reverse slot order - unfused.to(torch.bfloat16), # product rounded before the add - exact.to(torch.bfloat16), # float64 sum rounded once - bf16_sum, # bf16 accumulation - ): - with pytest.raises(AssertionError, match="differ bitwise"): - _compare(out, wrong) - - -# ── T = 0 and the loud rejections ───────────────────────────────────────── - - -def test_zero_tokens_returns_and_leaves_no_error(): - args, _, _ = _problem(_spec(8), seed=71) - empty = op( - permuted_input=args["permuted_input"], - output=torch.empty(0, H, dtype=torch.bfloat16, device=DEV), - expanded_idx_to_permuted_idx=args["expanded_idx_to_permuted_idx"][:0], - topk_scales=args["topk_scales"][:0], - ) - torch.cuda.synchronize() - assert tuple(empty.shape) == (0, H) - assert float((torch.ones(1000, device=DEV) + 1).sum()) == 2000.0 - out = op(**args) - _compare(out, op.reference(**args)) - - -def test_op_rejects_loudly(): - """Rejected by the op itself; none of these needs a guard.""" - tokens = 64 - args, _, _ = _problem(_spec(tokens), seed=81) - p, e, s = ( - args["permuted_input"], - args["expanded_idx_to_permuted_idx"], - args["topk_scales"], - ) - narrow = 2556 - cases = [ - ("int64 indices", dict(expanded_idx_to_permuted_idx=e.long()), "expected scalar type Int"), - ( - "hidden not a multiple of 8", - dict( - permuted_input=p[:, :narrow].contiguous(), - output=torch.empty(tokens, narrow, dtype=torch.bfloat16, device=DEV), - ), - "hidden_size must be divisible by 8", - ), - ( - "fp32 rows", - dict(permuted_input=p.float(), output=_output(tokens, torch.float32)), - "Unsupported input dtype", - ), - ("fp16 scales, bf16 rows", dict(topk_scales=s.half()), "Unsupported input dtype"), - ( - "fewer than T * top_k rows", - dict(permuted_input=p[: tokens * K - 1]), - r"greater than or equal to num_tokens \* top_k", - ), - ("scales rows", dict(topk_scales=s[:-1]), r"topk_scales.size\(0\) must be num_tokens"), - ( - "scales width", - dict(topk_scales=s[:, :9].contiguous()), - r"topk_scales.size\(1\) must be top_k", - ), - ( - "index rows", - dict(expanded_idx_to_permuted_idx=e[:-1]), - r"expanded_idx_to_permuted_idx.size\(0\) must be num_tokens", - ), - ( - "output width", - dict(output=torch.empty(tokens, H - 8, dtype=torch.bfloat16, device=DEV)), - r"output.size\(1\) must be hidden_size", - ), - ] - for tag, change, match in cases: - with pytest.raises(RuntimeError, match=match): - op(**dict(args, **change)) + assert _same_bits(out, eager), trial + del graph, e2p torch.cuda.synchronize() - out = op(**args) - _compare(out, op.reference(**args)) + assert stale_rows > 0, "no replay ran over a stale row: the stale-row claim is untested" # ── the guards ──────────────────────────────────────────────────────────── @@ -484,13 +252,12 @@ def _strided(t: torch.Tensor, pad: int) -> torch.Tensor: return wide[:, : t.shape[1]] -@pytest.mark.parametrize("name", ("permuted_input", "output")) -def test_noncontiguous_operand_answers_wrongly_and_is_refused(name): - """One operand the op reads and the one it writes; the guard checks all four - the same way.""" - args, _, _ = _problem(_spec(300), seed=61) +def test_noncontiguous_operand_answers_wrongly_and_is_refused(): + """The FC2 rows the op reads; the guard checks all four operands the same way.""" + name = "permuted_input" + args = _problem(_spec(300), seed=61) ref = op.reference(**args) - bad = dict(args, **{name: _strided(args[name], 64 if args[name].shape[1] == H else 6)}) + bad = dict(args, **{name: _strided(args[name], 64)}) out = op(**bad) with pytest.raises(AssertionError, match="differ bitwise"): _compare(out.contiguous(), ref) @@ -499,27 +266,3 @@ def test_noncontiguous_operand_answers_wrongly_and_is_refused(name): op(**bad) out = op(**dict(args, output=_output(300))) _compare(out, ref) - - -def test_output_dtype_mismatch_answers_wrongly_and_is_refused(): - """An fp32 output receives the bf16 result's bit patterns, packed into the - first half of its bytes. - """ - args, _, _ = _problem(_spec(300), seed=62) - ref = op.reference(**args) - out32 = op(**dict(args, output=_output(300, torch.float32))) - assert _same_bits(out32.view(torch.bfloat16).reshape(-1)[: 300 * H].view(300, H), ref) - with pytest.raises(AssertionError): - _compare(out32, ref) - with validating(op): - with pytest.raises(AssertionError, match="output dtype"): - op(**dict(args, output=_output(300, torch.float32))) - - -def test_tokens_past_int32_are_refused(): - """The guard reads only `output.size(0)` first, so a meta tensor stands in for - the 2^31-row output the op would silently leave unwritten.""" - args, _, _ = _problem(_spec(8), seed=63) - huge = torch.empty((2**31, 8), dtype=torch.bfloat16, device="meta") - with validating(op), pytest.raises(AssertionError, match=r">= 2\^31 rows"): - op(**dict(args, output=huge)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_renorm_moe_routing_op.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_renorm_moe_routing_op.py index fcf05d43e29e..5c7fe656dd28 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_renorm_moe_routing_op.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_renorm_moe_routing_op.py @@ -3,9 +3,8 @@ """GPU test for the renorm_moe_routing_op catalog entry. The entry owns the contract. This file turns a cell into logits and holds the -claims no cell carries: the return order, the tie-break against `torch.topk`, -the sibling op one flag away, the sweep of the old certified surface, the edge -behaviour, and the guards. +claims no cell carries: equivalence to HF's `norm_topk_prob` routing, and the +contiguity guard. """ import zlib @@ -95,211 +94,32 @@ def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves th ) -def test_old_certified_surface(): - """The rest of the surface the entry was certified over, through the entry's - own reference and gate: the token column in both output dtypes, every input - dtype against every output dtype, and the num_experts / topk sweep across - the kernel's power-of-two instantiation boundaries.""" - for tokens in (1, 2, 7, 8, 64, 512, 2048, 4096, 8192): - logits = _logits(tokens, 512, torch.bfloat16, seed=tokens) - for out_dtype in (torch.float32, torch.bfloat16): - _check(logits, 10, out_dtype) - base = _logits(256, 512, torch.float32, magnitude=2.0, seed=8) - for in_dtype in (torch.bfloat16, torch.float16, torch.float32): - for out_dtype in (torch.float32, torch.bfloat16, None): - _check(base.to(in_dtype), 10, out_dtype) - for num_experts in (1, 2, 8, 31, 32, 33, 64, 96, 128, 129, 256, 384, 511, 512): - _check(_logits(64, num_experts, torch.bfloat16, seed=num_experts), min(10, num_experts)) - logits = _logits(256, 512, torch.bfloat16, seed=9) - for topk in (1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 16): # topk 0: test_degenerate_shapes - _check(logits, topk, torch.bfloat16) - - -# ── the traps ───────────────────────────────────────────────────────────── - - -def test_return_order_is_ids_then_weights(): - logits = _logits(256, 512, torch.bfloat16, seed=1) - first, second = op(logits, 10, torch.bfloat16) - assert first.dtype == torch.int32 and second.dtype == torch.bfloat16 - top1 = torch.sort(logits, dim=-1, descending=True, stable=True).indices[:, 0] - assert torch.equal(first[:, 0], top1.to(torch.int32)) +# ── what a cell cannot carry ────────────────────────────────────────────── def test_equivalent_to_hf_norm_topk_prob(): """Softmax over all experts, top-k, divide by the selected sum: the same operator up to rounding (5.96e-08 measured); skipping the renormalization lands three orders of magnitude outside that.""" - for dtype in (torch.bfloat16, torch.float32): - logits = _logits(1024, 512, dtype, magnitude=2.0, seed=2) - _, weights = op(logits, 10, torch.float32) - hf_values, _ = torch.topk(torch.softmax(logits.float(), dim=-1), 10, dim=-1) - hf_weights = hf_values / hf_values.sum(-1, keepdim=True) - # ids can differ on tie rows; the sorted weight multiset cannot - torch.testing.assert_close( - weights.sort(-1, descending=True).values, - hf_weights.sort(-1, descending=True).values, - rtol=1e-5, - atol=1e-7, - ) - assert (weights - hf_values).abs().max().item() > 1e-2 - - -def test_distinct_from_default_moe_routing_op(): - """The sibling one flag away: same ids, weights that do not sum to 1.""" - logits = _logits(512, 512, torch.bfloat16, seed=3) - ids, weights = op(logits, 10, torch.float32) - other_ids, other_weights = torch.ops.trtllm.default_moe_routing_op(logits, 10, torch.float32) - assert torch.equal(ids, other_ids) - row_sums = other_weights.sum(-1) - assert float(row_sums.max()) < 0.99 - assert (weights - other_weights).abs().max().item() > 1e-2 + logits = _logits(1024, 512, torch.bfloat16, magnitude=2.0, seed=2) + _, weights = op(logits, 10, torch.float32) + hf_values, _ = torch.topk(torch.softmax(logits.float(), dim=-1), 10, dim=-1) + hf_weights = hf_values / hf_values.sum(-1, keepdim=True) + # ids can differ on tie rows; the sorted weight multiset cannot torch.testing.assert_close( - other_weights / row_sums.unsqueeze(-1), weights, rtol=1e-6, atol=1e-7 + weights.sort(-1, descending=True).values, + hf_weights.sort(-1, descending=True).values, + rtol=1e-5, + atol=1e-7, ) + assert (weights - hf_values).abs().max().item() > 1e-2 -def test_ties_break_toward_lower_expert_index(): - """Opposite to torch.topk, and common at 512 bf16 experts.""" - logits = torch.zeros(4, 64, dtype=torch.bfloat16, device=DEV) - logits[0, :] = 1.0 - logits[1, 3] = logits[1, 7] = 5.0 - logits[2, ::2] = 2.0 - logits[3, 0] = logits[3, 63] = 1.0 - ids, _ = _check(logits, 4, torch.float32) - expected = [[0, 1, 2, 3], [3, 7, 0, 1], [0, 2, 4, 6], [0, 63, 1, 2]] - assert ids.tolist() == expected - - dense = _logits(4096, 512, torch.bfloat16, seed=4) - ids, _ = _check(dense, 10, torch.float32) - torch_topk_ids = torch.topk(dense, 10, dim=-1).indices.to(torch.int32) - disagreeing = int((ids != torch_topk_ids).any(-1).sum()) - selected = torch.gather(dense, 1, ids.long()) - tied = int((selected[:, :-1] == selected[:, 1:]).any(-1).sum()) - print(f" {tied} of 4096 rows tie inside the top-10; torch.topk disagrees on {disagreeing}") - assert disagreeing > 0 and tied > dense.shape[0] // 4 - for tokens, num_experts, topk in ((512, 128, 8), (256, 32, 4)): - _check(_logits(tokens, num_experts, torch.bfloat16, "coarse", seed=5), topk, torch.float32) - - -def test_slot_order_and_weight_monotonicity(): - logits = _logits(512, 512, torch.bfloat16, seed=6) - ids, weights = _check(logits, 10, torch.float32) - selected = torch.gather(logits.float(), 1, ids.long()) - assert bool((selected[:, :-1] >= selected[:, 1:]).all()) - assert bool((weights[:, :-1] >= weights[:, 1:]).all()) - - -def test_output_dtype_defaults_to_fp32_and_casts_after_the_softmax(): - logits = _logits(1024, 512, torch.float32, magnitude=3.0, seed=7) - default_ids, default_weights = op(logits, 10) - assert default_weights.dtype == torch.float32 - explicit_ids, explicit_weights = op(logits, 10, torch.float32) - assert torch.equal(default_ids, explicit_ids) and torch.equal(default_weights, explicit_weights) - _, bf16_weights = _check(logits, 10, torch.bfloat16) - order = torch.sort(logits, dim=-1, descending=True, stable=True).indices[:, :10] - narrowed = torch.gather(logits, 1, order).to(torch.bfloat16).float() - e = torch.exp(narrowed - narrowed.max(-1, keepdim=True).values) - cast_first = (e / e.sum(-1, keepdim=True)).to(torch.bfloat16) - differing = int((bf16_weights != cast_first).sum()) - print(f" cast-before-softmax disagrees on {differing} of {bf16_weights.numel()}") - assert differing > bf16_weights.numel() // 10 - - -# ── edges ───────────────────────────────────────────────────────────────── - - -def test_degenerate_shapes(): - ids, weights = op(_logits(8, 512, torch.bfloat16), 0, torch.float32) - assert ids.shape == (8, 0) and weights.shape == (8, 0) - ids, weights = op(torch.empty(0, 512, dtype=torch.bfloat16, device=DEV), 10, torch.float32) - assert ids.shape == (0, 10) and weights.shape == (0, 10) - assert ids.dtype == torch.int32 and weights.dtype == torch.float32 - - -def test_input_not_mutated_and_deterministic(): - logits = _logits(4096, 512, torch.bfloat16, seed=10) - snapshot = logits.clone() - ids_a, weights_a = op(logits, 10, torch.bfloat16) - ids_b, weights_b = op(logits, 10, torch.bfloat16) - assert torch.equal(logits, snapshot) - assert torch.equal(ids_a, ids_b) and torch.equal(weights_a, weights_b) - assert weights_a.data_ptr() not in (weights_b.data_ptr(), logits.data_ptr()) - stream = torch.cuda.Stream() - with torch.cuda.stream(stream): - ids_c, weights_c = op(logits, 10, torch.bfloat16) - stream.synchronize() - assert torch.equal(ids_a, ids_c) and torch.equal(weights_a, weights_c) - - -def test_non_finite_logits_poison_the_whole_row(): - logits = _logits(3, 512, torch.float32) - logits[0, 5], logits[1, 7] = float("inf"), float("nan") - logits[2, :] = float("-inf") - ids, weights = op(logits, 10, torch.float32) - assert bool(torch.isnan(weights).all()) - assert int(ids[0, 0]) == 5 and int(ids[1, 0]) == 7 - # -inf is a soft mask: excluded while topk fits in the finite experts... - masked = _logits(4, 512, torch.float32, seed=1) - masked[:, :8] = float("-inf") - masked_ids, masked_weights = op(masked, 10, torch.float32) - assert bool(torch.isfinite(masked_weights).all()) and int(masked_ids.min()) >= 8 - # ...and at weight exactly 0 once it does not - starved = _logits(4, 512, torch.float32, seed=2) - starved[:, 7:] = float("-inf") - starved_ids, starved_weights = op(starved, 10, torch.float32) - assert bool(torch.isfinite(starved_weights).all()) - assert bool((starved_weights[:, 7:] == 0).all()) and bool((starved_weights[:, :7] > 0).all()) - assert int(starved_ids.max()) < 512 - - -def test_op_rejects_unsupported_domains(): - """Loud on the op's own; none of these needs a guard.""" - logits = _logits(8, 512, torch.bfloat16) - cases = [ - ("1D router_logits", lambda: op(logits[0], 6)), - ("3D router_logits", lambda: op(logits.unsqueeze(0), 6)), - ("fp64 router_logits", lambda: op(logits.double(), 6)), - ("int32 router_logits", lambda: op(logits.int(), 6)), - ("fp16 output_dtype", lambda: op(logits, 6, torch.float16)), - ("fp16 output_dtype, fp16 logits", lambda: op(logits.half(), 6, torch.float16)), - ("fp64 output_dtype", lambda: op(logits, 6, torch.float64)), - ("int32 output_dtype", lambda: op(logits, 6, torch.int32)), - ("topk 17", lambda: op(logits, 17)), - ("negative topk", lambda: op(logits, -1)), - ("num_experts 513", lambda: op(_logits(8, 513, torch.bfloat16), 6)), - ("cpu router_logits", lambda: op(logits.cpu(), 6)), - ] - for tag, call in cases: - with pytest.raises((RuntimeError, ValueError, NotImplementedError)): - call() - - -# ── the guards ──────────────────────────────────────────────────────────── - - -def test_unguarded_hazards_are_silent(): - """What `is_valid` exists for: the raw op answers both of these.""" - wide = _logits(8, 1024, torch.bfloat16, seed=12) - sliced = wide[:, :512] - dense = wide.reshape(-1)[: 8 * 512].reshape(8, 512).contiguous() - strided_ids, _ = op(sliced, 10, torch.float32) - assert torch.equal(strided_ids, op(dense, 10, torch.float32)[0]), ( - "the strided read is no longer a dense reinterpretation" - ) - assert not torch.equal(strided_ids, op.reference(sliced, 10, torch.float32)[0]) - narrow = _logits(4, 8, torch.bfloat16, seed=13) - padded_ids, padded_weights = op(narrow, 12, torch.float32) - assert bool((padded_ids[:, 8:] == 0).all()) and bool((padded_weights[:, 8:] == 0).all()) +# ── the guard ───────────────────────────────────────────────────────────── -def test_guards_refuse(): +def test_guard_refuses_a_column_slice(): + """What `is_valid` exists for: the raw op reads a slice as the dense buffer.""" wide = _logits(8, 1024, torch.bfloat16, seed=12) - with validating(op): - with pytest.raises(AssertionError, match="contiguous"): - op(wide[:, :512], 10) - with pytest.raises(AssertionError, match="contiguous"): - op(wide.t(), 6) - with pytest.raises(AssertionError, match="exceeds num_experts"): - op(_logits(4, 8, torch.bfloat16), 12) - _check(wide[:, :512].contiguous(), 10, torch.float32) + with validating(op), pytest.raises(AssertionError, match="contiguous"): + op(wide[:, :512], 10) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_flashinfer_gemma_rmsnorm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_flashinfer_gemma_rmsnorm.py index 0f2d08f82828..3d86b01b2f55 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_flashinfer_gemma_rmsnorm.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_flashinfer_gemma_rmsnorm.py @@ -57,65 +57,14 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res torch.testing.assert_close(out, op.reference(x, weight, eps)) -def test_weight_bias_is_one_not_zero() -> None: - """Discrimination: the scale is (1 + w), not w. - - Measured: 0 away from the (1 + w) reference, ~3.98 from the plain-w one. - """ - _, x, weight, eps = _build(dict(shape=(4, 5120), dtype=torch.bfloat16, eps=1e-6), seed=0) - out = op(x, weight, eps) - xf = x.float() - plain = (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps) * weight.float()).to(x.dtype) - gemma = op.reference(x, weight, eps) - - torch.testing.assert_close(out, gemma) - gemma_dev = (out.float() - gemma.float()).abs().max().item() - plain_dev = (out.float() - plain.float()).abs().max().item() - assert plain_dev > 100 * max(gemma_dev, 2.0**-8), (gemma_dev, plain_dev) - - def test_guard_refuses_fp8_before_the_kernel_answers() -> None: - """Both fp8 dtypes are in the DSL dtype table: accepted, and a tensor comes back.""" - for dtype in (torch.float8_e4m3fn, torch.float8_e5m2): - x = torch.randn(4, 2560, device="cuda").to(dtype) - w = torch.randn(2560, device="cuda").to(dtype) - unguarded = torch.ops.trtllm.flashinfer_gemma_rmsnorm(x, w, 1e-6) - assert unguarded.dtype is dtype and unguarded.shape == x.shape, ( - "the hazard this guard covers is gone" - ) - with pytest.raises(AssertionError, match="dequantize"): - with validating(op): - op(x, w, 1e-6) - - -def test_unsupported_dtypes_raise() -> None: - for dtype in (torch.float64, torch.int8, torch.uint8): - x = torch.ones(4, 2560, dtype=dtype, device="cuda") - w = torch.ones(2560, dtype=dtype, device="cuda") - with pytest.raises(KeyError): + """fp8 is in the DSL dtype table: accepted, and a tensor comes back.""" + x = torch.randn(4, 2560, device="cuda").to(torch.float8_e4m3fn) + w = torch.randn(2560, device="cuda").to(torch.float8_e4m3fn) + unguarded = torch.ops.trtllm.flashinfer_gemma_rmsnorm(x, w, 1e-6) + assert unguarded.dtype is x.dtype and unguarded.shape == x.shape, ( + "the hazard this guard covers is gone" + ) + with pytest.raises(AssertionError, match="dequantize"): + with validating(op): op(x, w, 1e-6) - - -def test_shape_dtype_and_device_mismatches_raise() -> None: - """All loud, which is why `is_valid` restates none of them.""" - w = torch.randn(2560, dtype=torch.bfloat16, device="cuda") - x = torch.randn(4, 2560, dtype=torch.bfloat16, device="cuda") - with pytest.raises(ValueError): - op(x, w.float(), 1e-6) - with pytest.raises(ValueError): - op(x, torch.randn(2048, dtype=torch.bfloat16, device="cuda"), 1e-6) - with pytest.raises(ValueError): - op(x, w.cpu(), 1e-6) - with pytest.raises(ValueError): - op(x.cpu(), w.cpu(), 1e-6) - with pytest.raises(ValueError): - op(torch.randn(2560, dtype=torch.bfloat16, device="cuda"), w, 1e-6) - # a strided normalized dim: "Mismatched mX.strides[1] ... expected to be 1" - with pytest.raises(ValueError, match="strides"): - op(torch.randn(8, 2560, 2, dtype=torch.bfloat16, device="cuda")[..., 0], w, 1e-6) - with pytest.raises(ValueError): - op( - torch.randn(2, 4, 8, 128, dtype=torch.bfloat16, device="cuda"), - torch.randn(128, dtype=torch.bfloat16, device="cuda"), - 1e-6, - ) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_layer_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_layer_norm.py index 921abba692ca..0e07d7975518 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_layer_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_layer_norm.py @@ -29,13 +29,6 @@ WEIGHT_LO, WEIGHT_HI = -1.210938, 17.5 BIAS_LO, BIAS_HI = -25.625, 14.875 -#: torch.testing's untouched defaults, restated for the discrimination test. -DEFAULT_TOL = { - torch.bfloat16: (1.6e-2, 1e-5), - torch.float16: (1e-3, 1e-5), - torch.float32: (1.3e-6, 1e-5), -} - def _gen(seed: int) -> torch.Generator: return torch.Generator(device=DEVICE).manual_seed(seed) @@ -69,157 +62,6 @@ def test_certified_cells(cell, certifying): # noqa: F811 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) @@ -227,32 +69,3 @@ def test_rejects_a_weight_whose_length_is_not_the_normalized_width(): weight, bias = _affine(gen) with pytest.raises(AssertionError): op(x, weight[:-1].contiguous(), bias, EPS) - with pytest.raises(AssertionError): - op(x, weight, bias[:-1].contiguous(), EPS) - - -def test_rejects_a_feature_dim_that_does_not_fit_the_fused_kernel(): - width = 65536 - x = torch.randn(2, width, device=DEVICE).to(torch.bfloat16) - weight = torch.ones(width, device=DEVICE, dtype=torch.bfloat16) - bias = torch.zeros(width, device=DEVICE, dtype=torch.bfloat16) - with pytest.raises(RuntimeError, match="doesn't support feature dim"): - op(x, weight, bias, EPS) - - -def test_accepts_an_fp8_activation_and_returns_fp8(): - """Measured, not certified: `empty_like(x)` makes the output e4m3, the - arithmetic is still fp32, and the store lands within one e4m3 ulp (2^-3 - relative, 2^-9 smallest subnormal) of the exact value.""" - gen = _gen(111) - x = torch.randn(64, WIDTH, generator=gen, device=DEVICE).to(torch.bfloat16) - weight, bias = _affine(gen) - x8 = x.to(torch.float8_e4m3fn) - - out = op(x8, weight, bias, EPS) - - assert out.dtype == torch.float8_e4m3fn - assert torch.isfinite(out.float()).all() - expected = _fp64(x8.float(), weight, bias, EPS) - deviation = (out.float().double() - expected).abs() - assert (deviation <= 2.0**-3 * expected.abs() + 2.0**-9).all(), deviation.max().item() diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_rms_norm_gated_token_major.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_rms_norm_gated_token_major.py index ca8bbbbb80ac..21569aeef4e0 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_rms_norm_gated_token_major.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_rms_norm_gated_token_major.py @@ -17,24 +17,12 @@ assert torch.cuda.is_available(), "rms_norm_gated_token_major requires a CUDA device" -def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: - """fp8 output: byte-exact, as measured. Otherwise torch's default band.""" - assert out.dtype == ref.dtype, (out.dtype, ref.dtype) - assert out.shape == ref.shape, (out.shape, ref.shape) - if out.dtype == torch.float8_e4m3fn: - assert torch.equal(out.view(torch.uint8), ref.view(torch.uint8)), ( - "fp8 output is byte-identical to the unfused reference" - ) - return - torch.testing.assert_close(out, ref) - - #: The target's concatenated GDN input projection: [q k v | z | a | b] = 10240 | 6144 | 48 | 48. _IN_PROJ_WIDTH, _Z_OFFSET = 16480, 10240 -#: The checkpoint's GDN epilogue geometry over a packed `z`, for the non-cell tests. -_PACKED = dict(op.CELLS[0].spec, tokens=17, z="packed") +#: The checkpoint's GDN epilogue geometry over a packed `z`, for the guard test. +_PACKED = dict(op.CELLS[0].spec, tokens=4, z="packed") def _seed(text: str) -> int: @@ -59,116 +47,34 @@ def randn(*shape): z = randn(t, heads, n) weight = randn(n) - fp8_scale = None - if "fp8_scale" in spec: - fp8_scale = torch.tensor([spec["fp8_scale"]], dtype=torch.float32, device="cuda") - return x, z, weight, spec["eps"], fp8_scale, spec["gate_activation"] + return x, z, weight, spec["eps"], None, spec["gate_activation"] @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name certifying(op) args = _build(cell.spec, seed=_seed(cell.why)) - x, z, weight, _, fp8_scale, _ = args - if cell.spec.get("z") == "in_proj" and cell.spec["tokens"] > 1: - assert z.stride() == (_IN_PROJ_WIDTH, cell.spec["width"], 1) and not z.is_contiguous() + x, z, weight = args[:3] + assert z.stride() == (_IN_PROJ_WIDTH, cell.spec["width"], 1) and not z.is_contiguous() before = [x.clone(), z.clone(), weight.clone()] out = op(*args) - assert out.shape == x.shape and out.is_contiguous() - assert out.dtype == (torch.float8_e4m3fn if fp8_scale is not None else x.dtype) + assert out.shape == x.shape and out.is_contiguous() and out.dtype == x.dtype for t, b in zip((x, z, weight), before): assert torch.equal(t, b), "inputs are read-only" - if fp8_scale is not None: - # ... and no byte is the e4m3 NaN pattern, so the byte gate compares numbers. - assert torch.isfinite(out.float()).all() - _compare(out, op.reference(*args)) - - -def test_fp8_byte_gate_discriminates() -> None: - """Each wrong variant of the reference misses the byte gate on most elements.""" - spec = dict(op.CELLS[-1].spec, fp8_scale=0.0234375) - 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), - } - 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 pass the sigmoid cells; the two gates differ.""" - x, z, weight, eps, _, _ = _build(_PACKED, seed=3) - assert not torch.equal( - op(x, z, weight, eps, None, "silu"), op(x, z, weight, eps, None, "sigmoid") - ) + torch.testing.assert_close(out, op.reference(*args)) def test_guard_refuses_a_weight_of_the_wrong_length() -> None: """Unguarded, a length-192 weight is silently truncated to weight[:128].""" - x, z, weight, eps, _, act = _build(dict(_PACKED, tokens=4), seed=12) + x, z, weight, eps, _, act = _build(_PACKED, seed=12) long_weight = torch.cat([weight, torch.randn(64, dtype=weight.dtype, device="cuda")]) assert torch.equal( torch.ops.trtllm.rms_norm_gated_token_major(x, z, long_weight, eps, None, act), torch.ops.trtllm.rms_norm_gated_token_major(x, z, weight, eps, None, act), ), "the hazard this guard covers is gone" - for bad in (long_weight, weight[:64]): - with pytest.raises(AssertionError, match="weight must be 1-D of length"): - with validating(op): - op(x, z, bad, eps, None, act) - - -def test_guard_refuses_a_malformed_fp8_scale() -> None: - """Unguarded, a two-element scale quantizes against element 0 and a bf16 - scale is accepted -- neither raises.""" - x, z, weight, eps, _, act = _build(dict(_PACKED, tokens=4), seed=12) - one = torch.tensor([0.5], dtype=torch.float32, device="cuda") - two = torch.tensor([0.5, 0.25], dtype=torch.float32, device="cuda") - half = torch.tensor([0.5], dtype=torch.bfloat16, device="cuda") - by_two = torch.ops.trtllm.rms_norm_gated_token_major(x, z, weight, eps, two, act) - by_one = torch.ops.trtllm.rms_norm_gated_token_major(x, z, weight, eps, one, act) - assert torch.equal(by_two.view(torch.uint8), by_one.view(torch.uint8)) - by_half = torch.ops.trtllm.rms_norm_gated_token_major(x, z, weight, eps, half, act) - assert by_half.dtype == torch.float8_e4m3fn - - with pytest.raises(AssertionError, match="fp8_scale must hold exactly one element"): - with validating(op): - op(x, z, weight, eps, two, act) - with pytest.raises(AssertionError, match="fp8_scale must be float32"): + with pytest.raises(AssertionError, match="weight must be 1-D of length"): with validating(op): - op(x, z, weight, eps, half, act) - - -def test_rejects_unsupported_gate_activation() -> None: - x, z, weight, eps, _, _ = _build(dict(_PACKED, 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(_PACKED, 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) - - -def test_prefill_launch_rejects_z_that_does_not_cover_x() -> None: - """From 32768 rows the op runs its own multi-row launch; a `z` one token short - still reaches the registered op's check instead of being read past its end.""" - x, z, weight, eps, _, act = _build(dict(_PACKED, tokens=683), seed=12) - assert x.shape[0] >= 32768, x.shape - with pytest.raises(AssertionError, match="does not match x shape"): - op(x, z[:-1], weight, eps, None, act) + op(x, z, long_weight, eps, None, act) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_gated_ngram_injection.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_gated_ngram_injection.py index c410d9e790ab..f93fdf147cd6 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_gated_ngram_injection.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_gated_ngram_injection.py @@ -1,13 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the gated_ngram_injection entry. +"""GPU test for the gated_ngram_injection entry: the cells (the sub-floor and +zero dots are a cell), the leading-dim guard and CUDA-graph replay.""" -Every mis-reading of this formula keeps the shapes and stays bounded by -`|value|`, so beyond the cells one test measures how far outside the entry's -band each wrong reading lands, and two pin the floor and the zero-dot gate. -""" - -import math import zlib import pytest @@ -132,119 +127,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res _compare(out, op.reference(*args)) -def test_flattened_call_matches_batched() -> None: - key, query, value = _make_operands(11, torch.bfloat16, (3, 17)) - batched = op(key, query, value, HC_COUNT) - flat = op(key.reshape(-1, WIDTH), query.reshape(-1, WIDTH), value.reshape(-1, HIDDEN), HC_COUNT) - assert torch.equal(flat, batched.reshape(-1, WIDTH)) - - -def test_output_is_bounded_by_the_value_it_gates() -> None: - """sigmoid is in (0, 1): no stream amplifies its value, to 4096x the operands.""" - key, query, value = _make_operands(40, torch.bfloat16, (256,)) - for scale in (1.0, 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,)) @@ -261,15 +143,6 @@ def test_guard_refuses_unequal_leading_dims() -> None: op(key, query, value[:1], HC_COUNT) -def test_malformed_widths_raise_rather_than_return() -> None: - """Why the width is not guarded: both mistakes are loud.""" - key, query, value = _make_operands(70, torch.bfloat16, (32,)) - with pytest.raises(RuntimeError): - op(key[:, : WIDTH - 8], query[:, : WIDTH - 8], value, HC_COUNT) - with pytest.raises(RuntimeError): - op(key, query, value, 3) - - def test_cuda_graph_capture_and_replay() -> None: key, query, value = _make_operands(100, torch.bfloat16, (8,)) warmup = torch.cuda.Stream() diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ngram_token_history.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ngram_token_history.py index 680aa9320e7c..0e16ce5bdedf 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ngram_token_history.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ngram_token_history.py @@ -5,9 +5,9 @@ Integer bookkeeping, gated bitwise -- window and the whole published pool -- against the entry's scalar reference, which builds each sequence's stream as a list and shares no structure with the masked per-shift select under test. -A wrong window is a window of valid ids, so beyond the cells: six near-miss -assemblies must each differ, hand-computed windows pin the lifecycle, chunked -prefill is checked step by step, and the two guards. +Beyond the cells: chunked prefill against one-shot (the cross-step carry +the forward relies on), the `query_start_loc` length guard and CUDA-graph +replay. """ import zlib @@ -116,28 +116,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res # ----------------------------------------------------------- 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) @@ -158,120 +136,6 @@ def call(pool, chunk, carry): assert torch.equal(pool, one_shot_pool) -# ----------------------------------------------------------- hand-computed - - -def _literal(ids, pool, slots, carries, starts): - window = op( - torch.tensor(ids, device=DEV), - pool, - torch.tensor(slots, dtype=torch.int32, device=DEV), - torch.tensor(carries, device=DEV), - torch.tensor(starts, dtype=torch.int32, device=DEV), - EOS, - NGRAM, - ) - return window.tolist() - - -def test_fresh_sequence_seeds_with_eos_not_the_pool() -> None: - pool = torch.full((4, CONTEXT_LEN), 1234, device=DEV) - window = _literal([11, 22, 33], pool, [2], [False], [0, 3]) - assert window == [[EOS, EOS, 11], [EOS, 11, 22], [11, 22, 33]] - assert pool[2].tolist() == [22, 33] - - -def test_carry_row_is_oldest_first_and_publishes_its_tail() -> None: - pool = torch.tensor([[70, 80]], device=DEV) - assert _literal([90, 100], pool, [0], [True], [0, 2]) == [[70, 80, 90], [80, 90, 100]] - assert pool[0].tolist() == [90, 100] - pool = torch.tensor([[70, 80]], device=DEV) - assert _literal([90], pool, [0], [True], [0, 1]) == [[70, 80, 90]] - assert pool[0].tolist() == [80, 90], "part carry, part new" - - -def test_zero_token_rows_keep_or_reseed() -> None: - pool = torch.tensor([[70, 80], [11, 12], [13, 14]], device=DEV) - window = _literal([90, 91], pool, [0, 1, 2], [True, True, False], [0, 2, 2, 2]) - assert window == [[70, 80, 90], [80, 90, 91]] - assert pool.tolist() == [[90, 91], [11, 12], [EOS, EOS]] - - -def test_sequence_boundary_is_not_crossed() -> None: - pool = torch.tensor([[70, 80], [170, 180]], device=DEV) - window = _literal([90, 91, 92, 190, 191], pool, [0, 1], [True, True], [0, 3, 5]) - assert window[3:] == [[170, 180, 190], [180, 190, 191]] - - -def test_in_stream_eos_is_passed_through_raw() -> None: - pool = torch.tensor([[70, 80]], device=DEV) - window = _literal([EOS, 91, 92], pool, [0], [True], [0, 3]) - assert window == [[70, 80, EOS], [80, EOS, 91], [EOS, 91, 92]] - - -# ----------------------------------------------------------- discrimination - - -def test_wrong_variants_are_separated() -> None: - """Six mis-assemblies, each in-vocabulary; each must differ in window or pool.""" - spec = dict( - lengths=(17, 1, 5, 1), - slots=(3, 0, 2, 1), - carries=(False, True, True, True), - pool_slots=6, - eos_density=0.08, - ) - args = _build(spec, seed=9001) - ids, pool, slots, carries, qsl = args[:5] - start = pool.clone() - window = op(*args) - published = pool - - def ref(**override): - kw = dict( - input_ids=ids, - context_pool=start.clone(), - state_indices=slots, - has_initial_states=carries, - query_start_loc=qsl, - eos_token_id=EOS, - ngram_size=NGRAM, - ) - kw.update(override) - return op.reference(**kw) - - right_window, right_pool = ref() - - def publishes_the_head(): - head_pool = start.clone() - starts = qsl.tolist() - for s, slot in enumerate(slots.tolist()): - if starts[s + 1] > starts[s]: - head_pool[slot] = right_window[starts[s + 1] - 1, :-1] - return right_window, head_pool - - variants = { - "newest_first": lambda: (right_window.flip(1), right_pool), - "zero_seed": lambda: ref(eos_token_id=0), - "ignores_carry_flags": lambda: ref(has_initial_states=torch.ones_like(carries)), - # One flat stream: each sequence's first tokens read its neighbour's tail. - "packs_across_boundaries": lambda: ( - ref( - state_indices=slots[:1], - has_initial_states=carries[:1], - query_start_loc=qsl[[0, -1]], - )[0], - right_pool, - ), - "rotated_pool_read": lambda: ref(context_pool=start.roll(1, dims=1)), - "publishes_the_head": publishes_the_head, - } - for name, build in variants.items(): - wrong_window, wrong_pool = build() - same = torch.equal(window, wrong_window) and torch.equal(published, wrong_pool) - assert not same, f"variant {name} is indistinguishable" - - # ----------------------------------------------------------- guards @@ -299,18 +163,6 @@ def _misses(args, bad_index, bad_value): return not (torch.equal(window, expected[0]) and torch.equal(bad[1], expected[1])) -def test_guard_refuses_carry_flags_of_the_wrong_length() -> None: - """Unguarded, one flag broadcasts: sequence 1 reads its stale slot, not the seed.""" - args = _two_sequences() - one_flag = torch.ones(1, dtype=torch.bool, device=DEV) - assert _misses(args, 3, one_flag), "the hazard this guard covers is gone" - bad = list(args) - bad[3] = one_flag - with pytest.raises(AssertionError, match="one flag per sequence"): - with validating(op): - op(*bad) - - def test_guard_refuses_a_short_query_start_loc() -> None: """Unguarded, a 2-entry vector for 2 sequences misattributes the tail.""" args = _two_sequences() @@ -323,35 +175,6 @@ def test_guard_refuses_a_short_query_start_loc() -> None: op(*bad) -@pytest.mark.parametrize( - "index,value,match", - [ - (1, lambda a: torch.zeros(4, CONTEXT_LEN + 1, dtype=torch.long, device=DEV), "must match"), - (3, lambda a: a[3].to(torch.int32), "where expected condition to be a boolean"), - (6, lambda a: 1, "must match the size of tensor"), - ], - ids=["pool_too_wide", "int_flags", "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 diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_gate_conv_inject_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_gate_conv_inject_norm.py index 3949c1e76bec..b2729a8e1ee6 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_gate_conv_inject_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_gate_conv_inject_norm.py @@ -5,8 +5,7 @@ Cells are gated against the float64 reference that rounds where the op rounds: the stream (`hc_out`) and the published conv state within one ulp of each element, the normed stream within two (see `_compare`). Beyond the cells: -agreement with the composition it replaces at the same band, separation from -plausible wrong formulas, and graph capture. +the operand-width guard and graph capture. """ import zlib @@ -14,15 +13,6 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_inject_norm import ( # noqa: E501 - qwen4_exp_hc_inject_norm, -) -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, -) -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, -) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.ple_gate_conv_inject_norm import ( # noqa: E501 CKPT_CLAMP_MIN, CKPT_DILATION, @@ -32,9 +22,6 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.ple_gate_conv_inject_norm import ( # noqa: E501 ple_gate_conv_inject_norm as op, ) -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, -) __extra_import_path__ = ["../../../.."] from _validating import ( # noqa: E402,F401 — path above; fixture by name @@ -162,98 +149,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by ) -def _grouped(x, w, c): - d = x.shape[1] // c - parts = torch.split(x, [d] * c, 1) - return torch.cat( - [flashinfer_gemma_rmsnorm(p, s, EPS) for s, p in zip(torch.split(w, d, 0), parts)], 1 - ) - - -def _old( - hc, - key, - value, - nk, - nq, - nc, - cw, - states, - slots, - qsl, - nx, - eps, - c, - clamp, - dil, - pad, - _p, - carries=None, -): - """The calls the decode / prefill PLE block and the site after it made before this entry.""" - gated = gated_ngram_injection(_grouped(key, nk, c), _grouped(hc, nq, c), value, c, clamp) - convolved = dilated_short_conv_paged( - x=torch.transpose(_grouped(gated, nc, c), 0, 1), - weight=cw, - bias=None, - conv_states=states, - query_start_loc=qsl, - cache_indices=slots, - has_initial_state=torch.ones_like(slots, dtype=torch.bool) if carries is None else carries, - dilation=dil, - silu_activation=True, - pad_slot_id=pad, - ) - hc_new = torch.add(hc, torch.add(gated, torch.transpose(convolved, 0, 1))) - _, normed = qwen4_exp_hc_inject_norm(hc_new, None, None, nx, eps, c) - return hc_new, normed, states - - -@pytest.mark.parametrize( - "lengths,carries", - [ - ((1,) * 128, None), - ((4,) * 8, None), - ((5000, 3000) + (1,) * 32, (False, True) + (True,) * 32), - ], - ids=["m128", "mtp4x8", "mixed2ctx32dec"], -) -def test_matches_the_composition_it_replaces(lengths, carries) -> None: - args = _build(dict(lengths=lengths, pool_slots=461, carries=carries), seed=len(lengths)) - old_args = list(args) - old_args[7] = args[7].clone() - want = _old(*old_args) - got = _run(args) - _compare(got, want) - for name, a, b in zip(("hc_out", "normed", "conv_states"), got, want): - worst, frac = _stats(a, b) - print( - f"\n {str(len(lengths)):>4} seqs {name:<12} worst {worst:.3g} ulp, differ {frac:.2e}", - end="", - ) - - -def test_wrong_variants_land_outside_the_band() -> None: - """Each mis-reading keeps every shape; each lands far outside the band.""" - args = _build(dict(lengths=(1,) * 16, pool_slots=461), seed=21) - want = op.reference(*args) - variants = { - "dilation 1": dict(index=14, value=1), - "clamp floor 1.0": dict(index=13, value=1.0), - "next norm = query norm": dict(index=10, value=args[4]), - "conv norm = key norm": dict(index=5, value=args[3]), - } - for name, change in variants.items(): - wrong = list(args) - wrong[change["index"]] = change["value"] - if name == "dilation 1": - wrong[7] = wrong[7][:, :, :3].contiguous() - got = _run(wrong) - worst = max(_ulp_ratio(got[i], want[i]).max().item() for i in (0, 1)) - print(f"\n {name:<26} {worst:10.4g} ulp", end="") - assert worst > 5.0, name - - def test_guard_refuses_a_short_value() -> None: args = _build(dict(lengths=(1, 1)), seed=4) args[2] = args[2][:, :-1] diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_ngram_embed.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_ngram_embed.py index 678e021561b9..c46ea883dc07 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_ngram_embed.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_ngram_embed.py @@ -5,8 +5,8 @@ Integer bookkeeping plus a one-rounding dequantize, so gated bitwise -- the embedding and the whole published pool -- against the entry's reference, which is the scalar references of `ngram_token_history` and `ngram_hash` and a -float64 gather-and-scale. Beyond the cells: bitwise agreement with the four -calls it replaces, and graph capture. +float64 gather-and-scale. Beyond the cells: the `query_start_loc` length guard +and graph capture. """ import zlib @@ -33,9 +33,6 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.ple_ngram_embed import ( # noqa: E501 ple_ngram_embed as op, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.quantization.dequantize_e4m3_per_tensor import ( # noqa: E501 - dequantize_e4m3_per_tensor, -) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.weights import ( # noqa: E501 _host_resident_empty, ) @@ -155,39 +152,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by assert args[1].data_ptr() != 0 and not torch.equal(args[1], pool_before) -def _old(ids, pool, slots, carries, qsl, table, scale, mult, vocab, offsets, eos, max_rows): - """The four calls the decode PLE block made before this entry.""" - seeded = torch.ones_like(slots, dtype=torch.bool) if carries is None else carries - window = ngram_token_history(ids, pool, slots, seeded, qsl, eos, mult.numel()) - rows = torch.reshape(ngram_hash(window, mult, vocab, offsets, eos, 1), [-1, vocab.numel()]) - gathered = torch.nn.functional.embedding(rows, table) - return torch.reshape(dequantize_e4m3_per_tensor(gathered, scale), [rows.shape[0], -1]) - - -@pytest.mark.parametrize( - "spec", - [ - dict(lengths=(1,) * 128, pool_slots=ENGINE_SLOTS, eos_density=0.05), - dict(lengths=(4,) * 32, pool_slots=ENGINE_SLOTS, eos_density=0.05), - dict( - lengths=(5000, 3000) + (1,) * 32, - carries=(False, True) + (True,) * 32, - pool_slots=ENGINE_SLOTS, - eos_density=0.05, - ), - ], - ids=["decode128", "mtp4x32", "mixed2ctx32dec"], -) -def test_matches_the_calls_it_replaces(spec) -> None: - args = _build(spec, seed=len(spec["lengths"])) - old_args = list(args) - old_args[1] = args[1].clone() - want = _old(*old_args) - got = op(*args) - assert torch.equal(got, want) - assert torch.equal(args[1], old_args[1]) - - def test_guard_refuses_short_query_start_loc() -> None: args = _build(dict(lengths=(1, 1, 1)), seed=3) args[4] = args[4][:-1] diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_bf16_to_fp32.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_bf16_to_fp32.py index 7196d321e733..3bf21aa13b26 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_bf16_to_fp32.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_bf16_to_fp32.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """GPU test for the bf16_to_fp32 entry: the cells bit for bit `x.float()`, graph capture, the -guards.""" +contiguity guard.""" import zlib @@ -61,9 +61,7 @@ def test_cuda_graph_capture_and_replay() -> None: assert _bitwise(out, x.float()) -def test_guards() -> None: +def test_guard_refuses_a_strided_input() -> None: x = _build(dict(shape=(8, 64)), 3) - with validating(op), pytest.raises(AssertionError): - op(x.float()) with validating(op), pytest.raises(AssertionError): op(x[:, ::2]) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_dequantize_e4m3_per_tensor.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_dequantize_e4m3_per_tensor.py index 0609b96981c0..ef52eab0be66 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_dequantize_e4m3_per_tensor.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_dequantize_e4m3_per_tensor.py @@ -3,9 +3,8 @@ """GPU test for the dequantize_e4m3_per_tensor catalog entry. The entry owns the contract. This file turns a cell into operands and holds -the claims no cell carries: the direction, the round trip through the -quantizer, degenerate and overflowing scales, graph capture, and the loud -rejections. +the claims no cell carries: graph capture with the scale read at replay, and +the loud rejection of an empty batch. """ import zlib @@ -52,11 +51,6 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: DEV = "cuda" -E4M3_MAX = 448.0 -# e4m3: a normal value's relative half-step and the smallest subnormal's -# absolute half-step, in units of the scale. -E4M3_REL_HALF_STEP = 2.0**-4 -E4M3_SUBNORMAL_HALF_STEP = 2.0**-10 _TABLE = {} @@ -79,10 +73,9 @@ def _input(spec, seed: int) -> torch.Tensor: _TABLE["table"] = _fp8(torch.randn(1 << 20, 160, device=DEV, generator=tg) * 40.0) idx = torch.randint(0, 1 << 20, shape[:2], device=DEV, generator=g) x = _TABLE["table"][idx] - elif source == "all_codes": - x = torch.arange(256, dtype=torch.uint8, device=DEV).view(shape).view(torch.float8_e4m3fn) else: - x = _fp8(torch.randn(*shape, device=DEV, generator=g) * 8.0) + assert source == "all_codes", source + x = torch.arange(256, dtype=torch.uint8, device=DEV).view(shape).view(torch.float8_e4m3fn) assert x.is_contiguous() and x.shape == shape return x @@ -112,55 +105,6 @@ def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves th # ── what a cell cannot carry ────────────────────────────────────────────── -def test_multiplies_rather_than_divides(): - x = _fp8(torch.tensor([[1.0, 2.0, -3.0, 0.5]], device=DEV)) - for value in (0.5, PLE_SCALE): - out = op(x, _scale(value, 2, torch.float32)) - assert torch.equal(_bits(out), _bits(x.float() * value)) - assert not torch.allclose(out, x.float() / value) - - -def test_top_of_the_range_at_the_checkpoint_scale(): - top = op(_fp8(torch.tensor([[E4M3_MAX, -E4M3_MAX]], device=DEV)), _scale(PLE_SCALE, 2)) - # 448 * PLE_SCALE = 0.08929443359375 exactly; bf16's nearest is this - assert top.float().flatten().tolist() == [0.08935546875, -0.08935546875] - - -def test_round_trip_through_the_static_quantizer(): - """The error floor is e4m3 itself; a reciprocated scale lands >100x out.""" - g = torch.Generator(device=DEV).manual_seed(15) - x = (torch.randn(256, 5120, device=DEV, generator=g) * 2.0).to(torch.bfloat16) - value = 0.05 - scales = _scale(value, 2, torch.float32) - q, _ = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor(x, scales) - back = op(q, scales) - _compare(back, op.reference(q, scales)) - bound = E4M3_REL_HALF_STEP * x.float().abs() + E4M3_SUBNORMAL_HALF_STEP * value - assert bool(((back - x.float()).abs() <= bound).all()) - wrong = op(q, _scale(1.0 / value, 2, torch.float32)) - assert float(((wrong - x.float()).abs() / bound)[x.float().abs() > 0.1].min()) > 100.0 - - -def test_degenerate_scale_values_are_accepted_and_silent(): - y = _fp8(torch.tensor([[1.0, -1.0, 0.0, E4M3_MAX]], device=DEV)) - out = op(y, _scale(0.0, 2)) - assert out.float().flatten().tolist() == [0.0, -0.0, 0.0, 0.0] - assert not bool(torch.isnan(out.float()).any()) - assert op(y, _scale(-2.0, 2)).float().flatten().tolist() == [-2.0, 2.0, -0.0, -896.0] - out = op(y, _scale(float("inf"), 2)).float().flatten() - assert out[0].item() == float("inf") and out[1].item() == float("-inf") - assert torch.isnan(out[2]) - assert bool(torch.isnan(op(y, _scale(float("nan"), 2)).float()).all()) - - -def test_overflow_goes_to_infinity_not_to_a_clamp(): - y = _fp8(torch.tensor([[1.0, E4M3_MAX]], device=DEV)) - out = op(y, _scale(1000.0, 2, torch.float16)).float().flatten() - assert out[0].item() == 1000.0 and out[1].item() == float("inf") - out = op(y, _scale(3.0e38, 2, torch.bfloat16)).float().flatten() - assert torch.isfinite(out[0]) and out[1].item() == float("inf") - - def test_cuda_graph_capture_and_replay(): """Both operands stay live inputs: the scale is read at replay.""" g = torch.Generator(device=DEV).manual_seed(16) @@ -190,48 +134,7 @@ def test_cuda_graph_capture_and_replay(): _compare(static_out, op.reference(fresh, _scale(0.5, 3))) -def test_deterministic_and_operands_untouched(): - g = torch.Generator(device=DEV).manual_seed(17) - x = _fp8(torch.randn(64, 16, 160, device=DEV, generator=g) * 8.0) - scales = _scale(PLE_SCALE, 3) - x_before, s_before = x.clone(), scales.clone() - out1, out2 = op(x, scales), op(x, scales) - assert torch.equal(_bits(out1), _bits(out2)) - assert torch.equal(x.view(torch.uint8), x_before.view(torch.uint8)) - assert torch.equal(_bits(scales), _bits(s_before)) - assert out1.data_ptr() not in (x.data_ptr(), scales.data_ptr()) - - -def test_rejects_unsupported(): - """Every metadata violation is loud, which is why the entry has no guard.""" - x = _fp8(torch.randn(4, 128, device=DEV)) - scales = _scale(PLE_SCALE, 2) - bf = dict(dtype=torch.bfloat16, device=DEV) - cases = { - "1-D input": lambda: op(x.reshape(-1), _scale(0.5, 1)), - "0-D input": lambda: op(_fp8(torch.tensor(1.0, device=DEV)), _scale(0.5, 0)), - "zero rows": lambda: op(_fp8(torch.randn(0, 128, device=DEV)), scales), - "zero cols": lambda: op(_fp8(torch.randn(4, 0, device=DEV)), scales), - "column slice": lambda: op(_fp8(torch.randn(8, 256, device=DEV))[:, :128], scales), - "transpose": lambda: op(_fp8(torch.randn(128, 8, device=DEV)).t(), scales), - "bfloat16 input": lambda: op(torch.randn(4, 128, **bf), scales), - "float8_e5m2 input": lambda: op( - torch.randn(4, 128, device=DEV).to(torch.float8_e5m2), scales - ), - "int8 input": lambda: op( - torch.randint(-8, 8, (4, 128), dtype=torch.int8, device=DEV), scales - ), - "float64 scales": lambda: op(x, _scale(0.5, 2, torch.float64)), - "cpu scales": lambda: op(x, scales.cpu()), - "cpu input": lambda: op(x.cpu(), scales), - "cpu input and scales": lambda: op(x.cpu(), scales.cpu()), - "scales rank below": lambda: op(x, _scale(0.5, 1)), - "scales rank above": lambda: op(x, _scale(0.5, 3)), - "0-D scales": lambda: op(x, _scale(0.5, 0)), - "scales dim not 1": lambda: op(x, torch.full((1, 2), 0.5, **bf)), - "scales leading dim not 1": lambda: op(x, torch.full((4, 1), 0.5, **bf)), - "empty scales": lambda: op(x, torch.empty((1, 0), **bf)), - } - for name, fn in cases.items(): - with pytest.raises((RuntimeError, NotImplementedError)): - fn() +def test_rejects_an_empty_batch(): + """The op raises on zero rows rather than launching; the caller branches first.""" + with pytest.raises((RuntimeError, NotImplementedError)): + op(_fp8(torch.randn(0, 128, device=DEV)), _scale(PLE_SCALE, 2)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_fp8_quantize_1x128.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_fp8_quantize_1x128.py index fdad8d2dfddc..caa087518151 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_fp8_quantize_1x128.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_fp8_quantize_1x128.py @@ -2,10 +2,10 @@ # SPDX-License-Identifier: Apache-2.0 """GPU test for the fp8_quantize_1x128 catalog entry. -The entry owns the contract. This file turns a cell into an input and holds -what a cell cannot carry: that every code disagreement is a rounding tie, the -`use_ue8m0` bit-exactness, the wrong-recipe discrimination, and the edge -behaviour the entry's `note` records. +The entry owns the contract. This file turns a cell into an input, holds +every code disagreement of a cell to a rounding tie, and adds what a cell +cannot carry: CUDA-graph capture and replay, and the loud rejection of a +strided input. """ import zlib @@ -14,8 +14,6 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.quantization.fp8_quantize_1x128 import ( # noqa: E501 - AMAX_FLOOR, - E4M3_MAX, SCALE_ULP_GATE, block_scales, ) @@ -79,10 +77,6 @@ def _compare( DEV = "cuda" -# e4m3: 3 mantissa bits, so a normal value carries at most a 2^-4 relative -# half-step; the smallest subnormal 2^-9 has a 2^-10 half-step, in scale units. -E4M3_REL_HALF_STEP = 2.0**-4 -E4M3_SUBNORMAL_HALF_STEP = 2.0**-10 # A code can only disagree when the exact product sits within the first # division's 2 ulp of an e4m3 rounding midpoint. TIE_REL_GATE = 2.0 * 2.0**-23 @@ -98,16 +92,6 @@ def _input(tokens: int, hidden: int, seed: int, magnitude: float = 3.0) -> torch return (torch.randn(tokens, hidden, device=DEV, generator=g) * magnitude).to(torch.bfloat16) -def _amax(x: torch.Tensor) -> torch.Tensor: - """Exact per-block amax `[T, H/128]` fp32, floored as the kernel floors it.""" - t, n = x.shape - return x.abs().reshape(t, n // 128, 128).amax(dim=2).float().clamp(min=AMAX_FLOOR) - - -def _ulp_f32(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: - return a.contiguous().view(torch.int32).long() - b.contiguous().view(torch.int32).long() - - def _assert_misses_are_ties(x: torch.Tensor, data: torch.Tensor, ref_data: torch.Tensor) -> float: """Every code disagreement is a rounding tie: the exact product sits within the scale's 2-ulp band of the midpoint between the two codes.""" @@ -141,205 +125,7 @@ def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves th print(f" {cell.why[:44]}: {100 * exact:.4f}% bit-exact, worst tie {rel:.3e}") -# ── the flag, the scale and the direction ───────────────────────────────── - - -def test_use_ue8m0_is_bit_exact_and_a_different_quantization(): - """Once the dequant scale is a power of two its reciprocal is exact, so the - approximate division cannot reach the result.""" - tokens, hidden = 17, 2560 - x = _input(tokens, hidden, seed=5 + tokens) - data, sf = op(x, True) - ref_data, ref_sf = op.reference(x, True) - assert torch.equal(sf.view(torch.int32), ref_sf.view(torch.int32)) - assert torch.equal(data.view(torch.uint8), ref_data.view(torch.uint8)) - assert bool((torch.frexp(sf)[0] == 0.5).all()) - d0, sf0 = op(x, False) - sf_changed = (sf != sf0).float().mean().item() - d_changed = (data.view(torch.uint8) != d0.view(torch.uint8)).float().mean().item() - print(f" ue8m0 T={tokens} H={hidden}: scales changed {sf_changed:.5f}, codes {d_changed:.5f}") - assert sf_changed > 0.9 and d_changed > 0.5 - assert bool((sf >= sf0 * (1 - 2**-20)).all()) and bool((sf <= sf0 * 2 * (1 + 2**-20)).all()) - - -def test_scale_is_amax_over_448_either_spelling(): - """Within 4 ulp of both `amax / 448` and the kernel's `1 / (448 / amax)`.""" - x = _input(1024, 2560, seed=21) - _, sf = op(x, False) - amax = _amax(x) - direct = (amax / E4M3_MAX).t() - two_step = (1.0 / (E4M3_MAX / amax)).t() - u_direct = _ulp_f32(sf, direct).abs().max().item() - u_two_step = _ulp_f32(sf, two_step).abs().max().item() - exact = int((sf.view(torch.int32) == two_step.contiguous().view(torch.int32)).sum()) - print( - f" max|ulp| vs amax/448 {u_direct}, vs 1/(448/amax) {u_two_step}; exact {exact}/{sf.numel()}" - ) - assert u_direct <= SCALE_ULP_GATE and u_two_step <= SCALE_ULP_GATE - # the amax itself is recoverable from the scale to within those ulps - rec = sf.t() * E4M3_MAX - assert bool(((rec - amax).abs() <= amax * (SCALE_ULP_GATE + 1) * 2**-23).all()) - - -def test_scale_multiplies_the_code(): - """Round trip bounded by e4m3's own constants; the reciprocal reading is - four orders of magnitude off.""" - x = _input(64, 2560, seed=2, magnitude=2.0) - data, sf = op(x, False) - cols = _expand(sf.t().contiguous(), 2560) - deq = data.float() * cols - bound = (E4M3_REL_HALF_STEP + (SCALE_ULP_GATE + 1) * 2**-23) * x.float().abs() - bound = bound + E4M3_SUBNORMAL_HALF_STEP * cols - assert bool(((deq - x.float()).abs() <= bound).all()) - - def rel_rms(a): - return ((a - x.float()).pow(2).mean().sqrt() / x.float().pow(2).mean().sqrt()).item() - - mul, div = rel_rms(deq), rel_rms(data.float() / cols) - print(f" multiply relRMS={mul:.5f} divide relRMS={div:.5g}") - assert mul < 0.05 and div > 100.0 - assert data.float().abs().max().item() == E4M3_MAX - - -def test_discrimination_against_wrong_recipes(): - """Each wrong reading of the amax map, scored against the same call.""" - t, n = 256, 2560 - x = _input(t, n, seed=99) - data, sf = op(x, False) - amax = _amax(x) - - def score(a_cols, amax_max=E4M3_MAX): - qs_cols = amax_max / a_cols - d = (x.float() * qs_cols).to(torch.float8_e4m3fn) - agree = (d.view(torch.uint8) == data.view(torch.uint8)).float().mean().item() - rel = (((1.0 / qs_cols)[:, ::128].t() - sf).abs() / sf).max().item() - return agree, rel - - mant, exp = torch.frexp(amax / E4M3_MAX) - pow2 = torch.ldexp(torch.ones_like(amax), torch.where(mant == 0.5, exp - 1, exp)) - variants = { - "correct": (_expand(amax, n), E4M3_MAX), - "ue8m0 power-of-two scale": (_expand(pow2 * E4M3_MAX, n), E4M3_MAX), - "one scale per tensor": (_expand(amax.max().expand_as(amax), n), E4M3_MAX), - "one scale per token row": ( - _expand(amax.amax(1, keepdim=True).expand_as(amax), n), - E4M3_MAX, - ), - "blocks down the token axis": ( - _amax(x.t().contiguous()).t().repeat_interleave(128, dim=0)[:t].contiguous(), - E4M3_MAX, - ), - "misaligned block boundaries": ( - torch.roll(_expand(_amax(torch.roll(x, -64, dims=1)), n), 64, dims=1).contiguous(), - E4M3_MAX, - ), - "e4m3 max taken as 240": (_expand(amax, n), 240.0), - "signed max instead of amax": ( - _expand(x.float().reshape(t, n // 128, 128).amax(dim=2).clamp(min=AMAX_FLOOR), n), - E4M3_MAX, - ), - } - for name, (a_cols, mx) in variants.items(): - agree, rel = score(a_cols, mx) - print(f" {name:34s} {agree:.5f} {rel:.4e}") - if name == "correct": - assert agree >= 0.995 and rel <= (SCALE_ULP_GATE + 1) * 2**-23 - else: - assert agree < 0.9 and rel > 1e-3, name - - -def test_transposed_scale_buffer_is_undetectable_at_the_square_shape(): - x = _input(20, 2560, seed=4, magnitude=1.0) - _, sf = op(x, False) - tr = sf.t().contiguous() - assert sf.shape == tr.shape == (20, 20) and tr.dtype == sf.dtype and tr.is_contiguous() - spread = ((tr - sf).abs() / sf).max().item() - print(f" worst relative difference between the buffer and its transpose: {spread:.4f}") - assert spread > 0.1 - - -def test_width_not_a_multiple_of_128_takes_its_own_tail_amax(): - x = torch.zeros(1, 144, dtype=torch.bfloat16, device=DEV) - x[0, 0], x[0, 130] = 4.0, 0.25 - _, sf = op(x, False) - assert abs(sf[0, 0].item() - 4.0 / E4M3_MAX) <= 4.0 / E4M3_MAX * 2**-20 - assert abs(sf[1, 0].item() - 0.25 / E4M3_MAX) <= 0.25 / E4M3_MAX * 2**-20 - - -# ── edges ───────────────────────────────────────────────────────────────── - - -def test_zero_block_takes_the_amax_floor(): - x = torch.zeros(4, 256, dtype=torch.bfloat16, device=DEV) - x[0, 0] = 1.0 - data, sf = op(x, False) - floor = AMAX_FLOOR / E4M3_MAX - assert abs(sf[0, 0].item() - 1.0 / E4M3_MAX) <= 1.0 / E4M3_MAX * 2**-20 - for b in range(2): - for t in range(4): - if (b, t) != (0, 0): - assert abs(sf[b, t].item() - floor) <= floor * 2**-20, (b, t, sf[b, t].item()) - assert bool((data[1:].float() == 0).all()) and bool((data[0, 1:].float() == 0).all()) - - -def test_nonfinite_input(): - """inf poisons its whole 128-block; NaN is confined to its own element.""" - x = torch.zeros(1, 128, dtype=torch.bfloat16, device=DEV) - x[0, 0], x[0, 1] = float("inf"), 1.0 - data, sf = op(x, False) - assert sf.flatten().tolist() == [float("inf")] - assert torch.isnan(data[0, 0].float()) and data[0, 1].float().item() == 0.0 - - x = torch.zeros(1, 128, dtype=torch.bfloat16, device=DEV) - x[0, 0], x[0, 1] = float("nan"), 1.0 - data, sf = op(x, False) - assert abs(sf[0, 0].item() - 1.0 / E4M3_MAX) <= 1.0 / E4M3_MAX * 2**-20 - assert torch.isnan(data[0, 0].float()) and data[0, 1].float().item() == E4M3_MAX - assert bool((data[0, 2:].float() == 0).all()) - - -def test_saturation_never_exceeds_the_block_max(): - x = torch.zeros(1, 256, dtype=torch.bfloat16, device=DEV) - x[0, 0], x[0, 1], x[0, 3], x[0, 128] = 1.0e4, 1.0, -1.0e4, 2.0**-12 - data, sf = op(x, False) - assert data[0, 0].float().item() == E4M3_MAX and data[0, 3].float().item() == -E4M3_MAX - assert data[0, 128].float().item() == E4M3_MAX - assert bool(torch.isfinite(data.float()).all()) - # bf16 rounds 1e4 to 9984, and that is the magnitude the scale carries - assert abs(sf[0, 0].item() - 9984.0 / E4M3_MAX) <= 9984.0 / E4M3_MAX * 2**-20 - - -def test_storage_padding_and_density(): - """`data` views a 4-row-padded allocation; both outputs are dense.""" - for t in (1, 2, 3, 5, 20, 64): - data, sf = op(_input(t, 2560, seed=t), False) - assert data.untyped_storage().size() == (t + 3) // 4 * 4 * 2560, t - assert data.is_contiguous() and data.stride() == (2560, 1) - assert sf.is_contiguous() and sf.untyped_storage().size() >= sf.numel() * 4 - assert sf.stride() == (t, 1) or sf.shape[0] == 1 - _, sf = op(_input(1, 128, seed=0), False) - assert sf.shape == (1, 1) and sf.is_contiguous() - - -def test_empty_input(): - data, sf = op(torch.empty(0, 2560, dtype=torch.bfloat16, device=DEV), False) - assert data.shape == (0, 2560) and data.dtype == torch.float8_e4m3fn - assert sf.shape == (20, 0) and sf.dtype == torch.float32 - - -def test_registered_fake_matches_eager_shapes(): - from torch._subclasses.fake_tensor import FakeTensorMode - - cells = [(1, 2560), (5, 2560), (20, 2560), (8192, 2560), (7, 144)] - eager = {} - for t, n in cells: - d, sf = op(_input(t, n, seed=0)) - eager[(t, n)] = ((tuple(d.shape), d.dtype), (tuple(sf.shape), sf.dtype)) - with FakeTensorMode(): - for t, n in cells: - d, sf = op(torch.empty(t, n, dtype=torch.bfloat16, device=DEV)) - got = ((tuple(d.shape), d.dtype), (tuple(sf.shape), sf.dtype)) - assert got == eager[(t, n)], (t, n, got, eager[(t, n)]) +# ── what a cell cannot carry ────────────────────────────────────────────── def test_cuda_graph_capture_and_replay(): @@ -368,31 +154,8 @@ def test_cuda_graph_capture_and_replay(): assert torch.equal(a_sf.view(torch.int32), e_sf.view(torch.int32)) -def test_input_untouched_and_deterministic(): - x = _input(128, 2560, seed=3) - before = x.clone() - d1, s1 = op(x, False) - d2, s2 = op(x, False) - assert torch.equal(x, before) - assert torch.equal(d1.view(torch.uint8), d2.view(torch.uint8)) - assert torch.equal(s1.view(torch.int32), s2.view(torch.int32)) - - -def test_rejects_unsupported(): - """Every precondition is the op's own and loud, which is why there is no guard.""" - bf = dict(dtype=torch.bfloat16, device=DEV) - cases = { - "float16 input": lambda: op(torch.randn(4, 128, dtype=torch.float16, device=DEV)), - "float32 input": lambda: op(torch.randn(4, 128, device=DEV)), - "float8 input": lambda: op(torch.randn(4, 128, **bf).to(torch.float8_e4m3fn)), - "cpu input": lambda: op(torch.randn(4, 128, dtype=torch.bfloat16)), - "1-D input": lambda: op(torch.randn(128, **bf)), - "3-D input": lambda: op(torch.randn(2, 4, 128, **bf)), - "width 120": lambda: op(torch.randn(4, 120, **bf)), - "width 2568": lambda: op(torch.randn(4, 2568, **bf)), - "transposed": lambda: op(torch.randn(128, 8, **bf).t()), - "column slice": lambda: op(torch.randn(8, 5120, **bf)[:, :2560]), - } - for name, fn in cases.items(): - with pytest.raises((RuntimeError, NotImplementedError)): - fn() +def test_rejects_a_strided_input(): + """The op's own precondition, loud, which is why there is no guard.""" + x = torch.randn(8, 5120, dtype=torch.bfloat16, device=DEV)[:, :2560] + with pytest.raises((RuntimeError, NotImplementedError)): + op(x) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_fwd.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_fwd.py index e4bcbccce9e2..5ca8d32dfed5 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_fwd.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_fwd.py @@ -2,7 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 """GPU test for the causal_conv1d_fwd entry.""" -import os import zlib from typing import Optional @@ -57,51 +56,29 @@ def _seed(text: str) -> int: def _build(spec, seed): """`__call__`'s positional arguments for one spec, freshly allocated.""" g = torch.Generator(device="cuda").manual_seed(seed) - dim, width = spec["dim"], spec.get("width", 4) - dtype, layout = spec.get("dtype", torch.bfloat16), spec.get("layout", "in_place") + dim, width, dtype = spec["dim"], 4, torch.bfloat16 def randn(*shape): return torch.randn(*shape, generator=g, device="cuda").to(dtype) - out = None - if "batch3d" in spec: - batch, seqlen = spec["batch3d"] - x = randn(batch, dim, seqlen) - qsl = None - else: - lengths = spec["lengths"] - total = sum(lengths) - starts = [0] - for length in lengths: - starts.append(starts[-1] + length) - qsl = torch.tensor(starts, dtype=torch.int32, device="cuda") - token_major = randn(total, dim) - if layout == "channel_last": - x = token_major.t() - assert x.stride() == (1, dim) - out = torch.zeros(total, dim, dtype=dtype, device="cuda").t() - else: - x = token_major.t().contiguous() - if layout == "out": - out = torch.zeros_like(x) + lengths = spec["lengths"] + starts = [0] + for length in lengths: + starts.append(starts[-1] + length) + qsl = torch.tensor(starts, dtype=torch.int32, device="cuda") + x = randn(sum(lengths), dim).t().contiguous() weight = randn(dim, width) - bias = randn(dim) if spec.get("bias", False) else None - pool = None if spec["pool_slots"] is None else randn(spec["pool_slots"], dim, width - 1) - slots = spec["slots"] - cache_indices = None if slots is None else torch.tensor(slots, dtype=torch.int32, device="cuda") - initial = spec["initial"] - his = None if initial is None else torch.tensor(initial, dtype=torch.bool, device="cuda") - silu = spec.get("silu", True) - pad = spec.get("pad_slot_id", PAD_SLOT_ID) - return [x, weight, bias, pool, qsl, cache_indices, his, silu, pad, out] + pool = randn(spec["pool_slots"], dim, width - 1) + cache_indices = torch.tensor(spec["slots"], dtype=torch.int32, device="cuda") + his = torch.tensor(spec["initial"], dtype=torch.bool, device="cuda") + return [x, weight, None, pool, qsl, cache_indices, his, True, PAD_SLOT_ID, None] def _clone(args): - """The mutable operands cloned (x, conv_states, out); strides preserved.""" + """The mutable operands cloned (x, conv_states).""" got = list(args) - for i in (0, 3, 9): - if got[i] is not None: - got[i] = got[i].clone() + for i in (0, 3): + got[i] = got[i].clone() return got @@ -118,9 +95,7 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res args = _build(cell.spec, seed=_seed(cell.why)) ref = op.reference(*args) returned, pool, got = _run(args) - assert returned is (got[0] if got[9] is None else got[9]), "returns the tensor it wrote" - if got[9] is not None: - assert torch.equal(got[0], args[0]), "x is read-only when out is given" + assert returned is got[0], "returns the tensor it wrote" _compare((returned, pool), ref) if cell.spec["pool_slots"] == ENGINE_SLOTS: # The whole-pool bitwise gate above already locks the unnamed slots; @@ -129,107 +104,28 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res assert not torch.equal(pool[slot], args[3][slot]), slot -def test_in_place_and_out_of_place_are_bitwise_identical() -> None: - spec = dict(dim=10240, lengths=(64, 9), pool_slots=4, slots=(1, 2), initial=(True, True)) - in_place, pool_a, _ = _run(_build(spec, seed=1)) - out_of_place, pool_b, _ = _run(_build(dict(spec, layout="out"), seed=1)) - assert torch.equal(in_place, out_of_place) and torch.equal(pool_a, pool_b) - - -def test_channel_last_and_channel_major_are_bitwise_identical() -> None: - """The two kernels agree bitwise under the default SiLU form only.""" - spec = dict( - dim=10240, lengths=(37, 1, 128), pool_slots=8, slots=(3, 0, 5), initial=(True, False, True) - ) - major, pool_a, _ = _run(_build(spec, seed=11)) - last, pool_b, _ = _run(_build(dict(spec, layout="channel_last"), seed=11)) - assert torch.equal(pool_a, pool_b) - if os.environ.get("TRTLLM_CAUSAL_CONV1D_SILU_MODE", "") in ("", "0", "1"): - assert torch.equal(major, last) - else: - torch.testing.assert_close(major, last) - - -def test_the_published_state_is_the_raw_input_tail() -> None: - """Discrimination: the state is the pre-convolution input, not the activation.""" - spec = dict(dim=512, lengths=(48,), pool_slots=1, slots=(0,), initial=(True,), layout="out") - args = _build(spec, seed=12) - returned, pool, _ = _run(args) - x = args[0] - assert torch.equal(pool[0], x[:, -3:]) - assert (pool[0].float() - returned[:, -3:].float()).abs().max().item() > 0.1 - - -def test_the_initial_state_is_read() -> None: - """Discrimination: flipping the carry moves the first W-1 columns, nothing else.""" - spec = dict(dim=512, lengths=(64,), pool_slots=1, slots=(0,)) - with_carry, _, _ = _run(_build(dict(spec, initial=(True,)), seed=13)) - without, _, _ = _run(_build(dict(spec, initial=(False,)), seed=13)) - assert (with_carry[:, :3].float() - without[:, :3].float()).abs().max().item() > 0.1 - assert torch.equal(with_carry[:, 3:], without[:, 3:]) - - -#: One refusal per family the op checks itself: a kernel/layout mismatch, a -#: metadata dtype, an operand dtype. -_REJECTED = ( - "channel_last_with_channel_major_out", - "query_start_loc_int64", - "conv_states_dtype", -) - - -@pytest.mark.parametrize("case", _REJECTED) -def test_op_rejects_loudly_what_is_valid_does_not_restate(case: str) -> None: - """What 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_with_channel_major_out": - args[0] = token_major.t() - args[9] = torch.empty(dim, tokens, dtype=torch.bfloat16, device="cuda") - elif case == "query_start_loc_int64": - args[4] = args[4].long() - elif case == "conv_states_dtype": - args[3], args[5] = pool.float(), slots - args[6] = torch.ones(2, dtype=torch.bool, device="cuda") - with pytest.raises(RuntimeError): - with validating(op): - op(*args) - - -def _guard_case(case: str): - """A fresh argument list exhibiting one silent violation.""" +def _token_strided_args(): + """A fresh argument list whose `x` has token stride 2: the channel-major + kernel ignores that stride and answers silently wrong.""" dim, tokens = 512, 64 spec = dict(dim=dim, lengths=(tokens,), pool_slots=1, slots=(0,), initial=(True,)) - if case == "width_five": - spec.update(width=5, pool_slots=None, slots=None, initial=None) args = _build(spec, seed=15) g = torch.Generator(device="cuda").manual_seed(16) - if case == "weight_dtype": - args[1] = args[1].float() - elif case == "token_stride": - wide = torch.randn(2 * tokens, dim, generator=g, device="cuda").to(torch.bfloat16) - args[0] = wide.t().contiguous()[:, ::2] - assert args[0].stride(1) == 2 - elif case == "state_stride": - wide = torch.randn(1, dim, 6, generator=g, device="cuda").to(torch.bfloat16) - args[3] = wide[:, :, ::2] + wide = torch.randn(2 * tokens, dim, generator=g, device="cuda").to(torch.bfloat16) + args[0] = wide.t().contiguous()[:, ::2] + assert args[0].stride(1) == 2 return args -@pytest.mark.parametrize("case", ["width_five", "weight_dtype", "token_stride", "state_stride"]) -def test_guard_refuses_what_the_op_answers_silently_wrong(case: str) -> None: +def test_guard_refuses_what_the_op_answers_silently_wrong() -> None: """First the hazard, unguarded: the op returns without raising and its answer misses the reference. Then the guard refuses the same call.""" - ref = op.reference(*_guard_case(case)) - unguarded = _guard_case(case) # strided views must not be cloned dense + ref = op.reference(*_token_strided_args()) + unguarded = _token_strided_args() # strided views must not be cloned dense returned = op(*unguarded) with pytest.raises(AssertionError): _compare((returned, unguarded[3]), ref) with pytest.raises(AssertionError): with validating(op): - op(*_guard_case(case)) + op(*_token_strided_args()) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_triton_update.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_triton_update.py index 53881dbb4a3c..7f8a993b6785 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_triton_update.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_triton_update.py @@ -8,7 +8,6 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.causal_conv1d_triton_update import ( # noqa: E501 - CKPT_DIM, ENGINE_SLOTS, ) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.causal_conv1d_triton_update import ( # noqa: E501 @@ -49,67 +48,43 @@ def _compare(out: tuple, ref: tuple) -> None: def _build(spec, seed): """`__call__`'s positional arguments for one spec, freshly allocated.""" g = torch.Generator(device="cuda").manual_seed(seed) - dim, batch = spec["dim"], spec["batch"] - width, state_len = spec.get("width", 4), spec.get("state_len", 3) - seqlen, dtype = spec.get("seqlen", 2), spec.get("dtype", torch.bfloat16) - x_form = spec.get("x_form", "token_major") + dim, batch, seqlen, width = spec["dim"], spec["batch"], spec["seqlen"], 4 def randn(*shape): - return torch.randn(*shape, generator=g, device="cuda").to(dtype) - - if x_form == "2d": - x = randn(batch, dim) - elif x_form == "contiguous": - x = randn(batch, dim, seqlen) - elif x_form == "in_proj": - # The call site under the one input GEMM: the q|k|v columns of an - # [n * p, C + 6240] projection, reshaped [n, p, C] and seen as [n, C, p]. - wide = dim + 6144 + 2 * 48 - x = randn(batch * seqlen, wide)[:, :dim].reshape(batch, seqlen, dim).transpose(1, 2) - assert x.stride() == (seqlen * wide, 1, wide) - else: # the call site: an [n, p, C] block seen as [n, C, p] - x = randn(batch, seqlen, dim).transpose(1, 2) - conv_state = randn(spec["pool_slots"], dim, state_len) + return torch.randn(*shape, generator=g, device="cuda").to(torch.bfloat16) + + # The call site under the one input GEMM: the q|k|v columns of an + # [n * p, C + 6240] projection, reshaped [n, p, C] and seen as [n, C, p]. + wide = dim + 6144 + 2 * 48 + x = randn(batch * seqlen, wide)[:, :dim].reshape(batch, seqlen, dim).transpose(1, 2) + assert x.stride() == (seqlen * wide, 1, wide) + conv_state = randn(spec["pool_slots"], dim, width - 1) weight = randn(dim, width) - bias = randn(dim) if spec.get("bias", False) else None - slots = spec["slots"] - if slots == "random": - indices = torch.randperm(spec["pool_slots"], generator=g, device="cuda")[:batch] - indices = indices.to(torch.int32) - elif slots is None: - indices = None - else: - indices = torch.tensor(slots, dtype=torch.int32, device="cuda") - accepted = spec.get("accepted") - if accepted is not None: - accepted = torch.tensor(accepted, dtype=torch.int32, device="cuda") - rows = spec.get("snapshots", 0) - window = rows_idx = None - if rows: - steps = spec.get("snapshot_steps", 1 if x_form == "2d" else seqlen) - window = torch.zeros(rows, steps, dim, width - 1, dtype=dtype, device="cuda") - rows_idx = torch.arange(batch, dtype=torch.int32, device="cuda") + indices = torch.tensor(spec["slots"], dtype=torch.int32, device="cuda") + steps = spec.get("snapshot_steps", seqlen) + window = torch.zeros( + spec["snapshots"], steps, dim, width - 1, dtype=torch.bfloat16, device="cuda" + ) + rows_idx = torch.arange(batch, dtype=torch.int32, device="cuda") return [ x, conv_state, weight, - bias, - spec.get("activation", "silu"), + None, + "silu", indices, - spec.get("pad_slot_id", PAD_SLOT_ID), - accepted, + PAD_SLOT_ID, + None, window, rows_idx, - spec.get("launch_dependent_kernels", False), + False, ] def _run(args): """Drive the op on copies of the two buffers; return the triple compare takes.""" got = list(args) - got[1] = args[1].clone() - if args[8] is not None: - got[8] = args[8].clone() + got[1], got[8] = args[1].clone(), args[8].clone() returned = op(*got) return returned, got[1], got[8] @@ -129,98 +104,13 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res assert not torch.equal(got[1][slot], args[1][slot]), slot -def test_the_last_snapshot_is_what_the_pool_holds() -> None: - """Independent of the reference: snapshot `seqlen - 1` equals the advanced slot.""" - spec = dict(dim=512, batch=6, seqlen=4, pool_slots=16, slots="random", snapshots=6) - args = _build(spec, seed=1) - _, pool, snapshots = _run(args) - for b, slot in enumerate(args[5].tolist()): - assert torch.equal(snapshots[b, 3], pool[slot]), b - - -def test_fp32_products_are_the_wrong_reference() -> None: - """Discrimination: the gate is default-tolerance tight, and the natural - wrong guess -- multiplying in fp32 -- lands outside it. - - The rival is `reference` on fp32 copies of the operands, rounded once at - the end. Measured at 10240 channels: the faithful reference differs on - 2.4e-5 of elements (one-ulp `exp` cases); the rival on 0.37 of them. - """ - spec = dict(dim=CKPT_DIM, batch=4, pool_slots=128, slots="random") - args = _build(spec, seed=0) - got = _run(args) - faithful = op.reference(*args) - _compare(got, faithful) - fp32_args = list(args) - for i in (0, 1, 2): - fp32_args[i] = args[i].float() - rival_out = op.reference(*fp32_args)[0].to(args[0].dtype) - rival = (rival_out, faithful[1], faithful[2]) - with pytest.raises(AssertionError): - _compare(got, rival) - faithful_diff = (got[0] != faithful[0]).float().mean().item() - rival_diff = (got[0] != rival_out).float().mean().item() - print(f"\nelements differing: faithful {faithful_diff:.3e}, rival {rival_diff:.3e}") - assert faithful_diff < 1e-3 and rival_diff > 0.3, (faithful_diff, rival_diff) - - -def test_token_major_and_contiguous_x_are_bitwise_identical() -> None: - args = _build(dict(dim=512, batch=3, seqlen=4, pool_slots=16, slots=(2, 5, 1)), seed=11) - assert not args[0].is_contiguous() - strided = _run(args) - packed = _run([args[0].contiguous()] + args[1:]) - assert torch.equal(strided[0], packed[0]) and torch.equal(strided[1], packed[1]) - - -def test_launch_dependent_kernels_is_bitwise_inert() -> None: - spec = dict(dim=512, batch=4, pool_slots=16, slots="random", snapshots=4) - off = _run(_build(spec, seed=12)) - on = _run(_build(dict(spec, launch_dependent_kernels=True), seed=12)) - assert all(torch.equal(a, b) for a, b in zip(off, on)) - - -def test_width_five_is_silently_wrong_and_refused() -> None: - """Unguarded, width 5 runs and misses the reference; the guard refuses it.""" - spec = dict(dim=512, batch=2, width=5, state_len=4, pool_slots=8, slots=(1, 3)) - args = _build(spec, seed=13) - ref = op.reference(*args) - with pytest.raises(AssertionError): - _compare(_run(args), ref) - with pytest.raises(AssertionError, match="width must be 2, 3 or 4"): - with validating(op): - op(*_build(spec, seed=13)) - - -def test_a_short_pool_writes_into_the_next_slot_and_is_refused() -> None: - """Unguarded, a 2-column pool at W 4 writes past its slot into an unnamed - neighbour; the guard refuses it.""" - spec = dict(dim=512, batch=2, state_len=2, pool_slots=8, slots=(1, 4)) - args = _build(spec, seed=14) - _, pool, _ = _run(args) - unnamed = [s for s in range(8) if s not in (1, 4)] - assert not torch.equal(pool[unnamed], args[1][unnamed]), "the hazard this guard covers is gone" - with pytest.raises(AssertionError, match="must hold at least 3 columns"): - with validating(op): - op(*_build(spec, seed=14)) - - -@pytest.mark.parametrize( - "case,match", - [ - # The pool-width guard; the weight shares it. - ("narrow_pool", "must cover all 512 channels"), - ("short_snapshot_buffer", "intermediate_conv_window must be at least"), - ], -) -def test_guard_refuses_undersized_buffers(case, 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) +def test_guard_refuses_undersized_buffers() -> None: + """A snapshot buffer shallower than seqlen is accepted by the kernel and + written past the row it was given (measured); the guard refuses it + before the launch.""" + spec = dict(dim=512, batch=2, seqlen=2, pool_slots=8, slots=(0, 2), snapshots=2) args = _build(spec, seed=15) - if case == "narrow_pool": - args[1] = torch.randn(8, 256, 3, dtype=torch.bfloat16, device="cuda") - elif case == "short_snapshot_buffer": - args[8] = torch.zeros(2, 1, 512, 3, dtype=torch.bfloat16, device="cuda") - with pytest.raises(AssertionError, match=match): + args[8] = torch.zeros(2, 1, 512, 3, dtype=torch.bfloat16, device="cuda") + with pytest.raises(AssertionError, match="intermediate_conv_window must be at least"): with validating(op): op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_update.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_update.py index baa1f27bcfaa..bbd3b86e9575 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_update.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_update.py @@ -39,80 +39,48 @@ def _compare( PAD_SLOT_ID = -1 -#: Rows of prefill ahead of the decode rows in the mixed-batch cell. +#: Rows of prefill ahead of the decode rows in the in_proj output. _PREFILL_ROWS = 37 #: Columns past q|k|v in the target's one GDN input-projection output: z | a | b. _IN_PROJ_TAIL = 6144 + 2 * 48 def _build(spec, seed): - """`__call__`'s arguments for one spec, plus the buffers some layouts view.""" + """`__call__`'s arguments for one spec, plus the buffers x and the pool view.""" g = torch.Generator(device="cuda").manual_seed(seed) - dim, batch = spec["dim"], spec["batch"] - width, state_len, seqlen = spec.get("width", 4), spec.get("state_len", 3), spec.get("seqlen", 1) - dtype, x_form = spec.get("dtype", torch.bfloat16), spec.get("x_form", "3d") + dim, batch, layers = spec["dim"], spec["batch"], spec.get("pool_layers", 1) def randn(*shape): - return torch.randn(*shape, generator=g, device="cuda").to(dtype) - - extras = {} - if x_form == "2d": - x = randn(batch, dim) - elif x_form == "2d_row_slice": - extras["buf"] = randn(_PREFILL_ROWS + batch, dim) - x = extras["buf"][_PREFILL_ROWS:] - elif x_form == "2d_in_proj": - extras["buf"] = randn(_PREFILL_ROWS + batch, dim + _IN_PROJ_TAIL) - x = extras["buf"][_PREFILL_ROWS:, :dim] - assert x.stride() == (dim + _IN_PROJ_TAIL, 1) - elif x_form == "token_major": - x = randn(batch, seqlen, dim).transpose(1, 2) - assert x.stride() == (seqlen * dim, 1, dim) - else: - x = randn(batch, dim, seqlen) - if spec.get("strided_state", False): - extras["wide"] = randn(spec["pool_slots"], dim, 2 * state_len) - conv_state = extras["wide"][:, :, ::2] - else: - conv_state = randn(spec["pool_slots"], dim, state_len) - weight = randn(dim, width) - bias = randn(dim) if spec.get("bias", False) else None + return torch.randn(*shape, generator=g, device="cuda").to(torch.bfloat16) + + buf = randn(_PREFILL_ROWS + batch, dim + _IN_PROJ_TAIL) + x = buf[_PREFILL_ROWS:, :dim] + assert x.stride() == (dim + _IN_PROJ_TAIL, 1) + pool_buf = randn(spec["pool_slots"], layers, dim, 3) + conv_state = pool_buf[:, layers - 1] + weight = randn(dim, 4) slots = spec["slots"] if slots == "random": indices = torch.randperm(spec["pool_slots"], generator=g, device="cuda")[:batch] indices = indices.to(torch.int32) - elif slots is None: - indices = None else: indices = torch.tensor(slots, dtype=torch.int32, device="cuda") - cache_seqlens = spec.get("cache_seqlens") - if cache_seqlens is not None: - cache_seqlens = torch.tensor(cache_seqlens, dtype=torch.int32, device="cuda") - silu, pad = spec.get("silu", True), spec.get("pad_slot_id", PAD_SLOT_ID) - return [x, conv_state, weight, bias, silu, cache_seqlens, indices, pad], extras + args = [x, conv_state, weight, None, True, None, indices, PAD_SLOT_ID] + return args, dict(buf=buf, pool_buf=pool_buf) def _run(args, extras): """Drive the op on copies of x and the pool, through the same views.""" got = list(args) - if "buf" in extras: - got_buf = extras["buf"].clone() - got[0] = got_buf[_PREFILL_ROWS:, : args[0].shape[1]] - else: - got[0] = args[0].clone() - if "wide" in extras: - got_wide = extras["wide"].clone() - got[1] = got_wide[:, :, ::2] - else: - got[1] = args[1].clone() + got_buf, got_pool = extras["buf"].clone(), extras["pool_buf"].clone() + got[0] = got_buf[_PREFILL_ROWS:, : args[0].shape[1]] + got[1] = got_pool[:, got_pool.shape[1] - 1] returned = op(*got) assert returned is got[0], "returns the x it was handed" - if "buf" in extras: - assert torch.equal(got_buf[:_PREFILL_ROWS], extras["buf"][:_PREFILL_ROWS]) - dim = args[0].shape[1] - assert torch.equal(got_buf[:, dim:], extras["buf"][:, dim:]), "columns past x written" - if "wide" in extras: - assert torch.equal(got_wide[:, :, 1::2], extras["wide"][:, :, 1::2]) + assert torch.equal(got_buf[:_PREFILL_ROWS], extras["buf"][:_PREFILL_ROWS]) + dim = args[0].shape[1] + assert torch.equal(got_buf[:, dim:], extras["buf"][:, dim:]), "columns past x written" + assert torch.equal(got_pool[:, :-1], extras["pool_buf"][:, :-1]), "another layer written" return returned, got[1] @@ -120,71 +88,24 @@ def _run(args, extras): def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name certifying(op) args, extras = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) - weight, bias = args[2].clone(), None if args[3] is None else args[3].clone() + weight = args[2].clone() ref = op.reference(*args) returned, state = _run(args, extras) assert torch.equal(args[2], weight), "weight is read-only" - assert bias is None or torch.equal(args[3], bias), "bias is read-only" _compare((returned, state), ref) if cell.spec["pool_slots"] == ENGINE_SLOTS: # Non-vacuity for the large pool: every named slot moved. for slot in cell.spec["slots"]: - assert not torch.equal(state[slot], args[1][slot]), slot - - -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] - - -#: One metadata violation per family the op checks: the conv width, the slot -#: indices, the operand dtypes. -_REJECTIONS = { - "width_5": ( - "causal_conv1d only supports width between 2 and 4", - dict(dim=128, width=5, state_len=4), - ), - "indices_int64": ("conv_state_indices.scalar_type", None), - "weight_dtype": ("weight type must equal to input type", None), -} + if slot != PAD_SLOT_ID: + assert not torch.equal(state[slot], args[1][slot]), slot -@pytest.mark.parametrize("case", sorted(_REJECTIONS)) -def test_op_rejects_metadata_violations_loudly(case: str) -> None: +def test_op_rejects_metadata_violations_loudly() -> None: """Why the entry has no `is_valid`: the op raises on metadata violations - itself (one probe per family here).""" - fragment, overrides = _REJECTIONS[case] - args = _small(**(overrides or {})) - if case == "indices_int64": - args[6] = args[6].long() - elif case == "weight_dtype": - args[2] = args[2].float() - with pytest.raises(RuntimeError, match=fragment.replace("(", r"\(").replace(")", r"\)")): + itself (probed here at int64 slot indices).""" + spec = dict(dim=512, batch=4, pool_slots=16, slots=(0, 1, 2, 3)) + args = _build(spec, seed=14)[0] + args[6] = args[6].long() + with pytest.raises(RuntimeError, match="conv_state_indices.scalar_type"): with validating(op): op(*args) - - -def test_decode_rows_match_the_trtllm_kernel_bitwise() -> None: - """From 64 rows the decode shape runs the CuTe kernel: at that boundary x and - the whole pool are bitwise what `torch.ops.trtllm.causal_conv1d_update` leaves, - a pad row untouched. - x is the q|k|v column view of an in_proj output; the pool a per-layer view of a - two-layer buffer, as the engine hands both over.""" - rows = 64 - g = torch.Generator(device="cuda").manual_seed(rows) - dim = 10240 - proj = torch.randn(rows, dim + _IN_PROJ_TAIL, generator=g, device="cuda").to(torch.bfloat16) - pool = torch.randn(ENGINE_SLOTS, 2, dim, 3, generator=g, device="cuda").to(torch.bfloat16)[:, 1] - weight = (torch.randn(dim, 4, generator=g, device="cuda") * 0.5).to(torch.bfloat16) - slots = (torch.randperm(ENGINE_SLOTS, generator=g, device="cuda")[:rows]).to(torch.int32) - slots[rows // 3] = PAD_SLOT_ID - x_ref, pool_ref = proj.clone(), pool.clone() - torch.ops.trtllm.causal_conv1d_update( - x_ref[:, :dim].unsqueeze(-1), pool_ref, weight, None, True, None, slots, PAD_SLOT_ID - ) - x_got, pool_got = proj.clone(), pool.clone() - op(x_got[:, :dim], pool_got, weight, None, True, None, slots, PAD_SLOT_ID) - assert torch.equal(x_got, x_ref), "x differs from the TRT-LLM kernel" - assert torch.equal(pool_got, pool_ref), "the pool differs from the TRT-LLM kernel" - assert torch.equal(x_got[rows // 3], proj[rows // 3]), "the pad row was written" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_chunk_gated_delta_rule.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_chunk_gated_delta_rule.py index 0c389d3a15b0..6a8cf2fa2286 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_chunk_gated_delta_rule.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_chunk_gated_delta_rule.py @@ -14,12 +14,12 @@ import sys import zlib from pathlib import Path -from typing import Optional import pytest import torch from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.chunk_gated_delta_rule import ( # noqa: E501 + ENGINE_SLOTS, HEAD_K_DIM, HEAD_V_DIM, NUM_K_HEADS, @@ -50,7 +50,7 @@ def _compare(out: tuple, ref: tuple) -> None: on states, about 4x and 6x the largest deviations measured warm over every cell against this reference (`o` 4.88e-4, one bf16 ulp, on magnitudes up to 7.9e-2; states 3.57e-3 on up to 8.3e-1). - Discriminated: seven wrong references (no gating, beta = 1, q/k swapped, + Discriminated when measured: seven wrong references (no gating, beta = 1, q/k swapped, initial state dropped, scale = 1, key heads reversed, `[K, V]` transposed state) all fail the `o` gate, the closest 11x outside it. """ @@ -66,7 +66,6 @@ def _compare(out: tuple, ref: tuple) -> None: DEV = "cuda" -SCALE = HEAD_K_DIM**-0.5 def _l2norm(x: torch.Tensor) -> torch.Tensor: @@ -74,38 +73,27 @@ def _l2norm(x: torch.Tensor) -> torch.Tensor: return xf / torch.sqrt(xf.pow(2).sum(-1, keepdim=True) + 1e-6) -def _operands( - seqlens: Optional[tuple], - seed: int, - batch: int = 1, - num_tokens: Optional[int] = None, - normalize_qk: bool = True, -): +def _operands(seqlens: tuple, seed: int): """Operands in the domain the GDN prefill path produces: unit-norm q/k, `g = -exp(A_log) * softplus(a + dt_bias) < 0`, `beta` a sigmoid. The delta rule is only contractive there.""" gen = torch.Generator(device=DEV).manual_seed(seed) - if seqlens is None: - cu_seqlens, total = None, num_tokens - else: - total = sum(seqlens) - cu_seqlens = torch.tensor( - [0] + torch.tensor(seqlens).cumsum(0).tolist(), dtype=torch.long, device=DEV - ) - shape_qk = (batch, total, NUM_K_HEADS, HEAD_K_DIM) + total = sum(seqlens) + cu_seqlens = torch.tensor( + [0] + torch.tensor(seqlens).cumsum(0).tolist(), dtype=torch.long, device=DEV + ) + shape_qk = (1, total, NUM_K_HEADS, HEAD_K_DIM) q = torch.randn(shape_qk, generator=gen, device=DEV) k = torch.randn(shape_qk, generator=gen, device=DEV) - if normalize_qk: - q, k = _l2norm(q), _l2norm(k) - q, k = q.to(torch.bfloat16), k.to(torch.bfloat16) + q, k = _l2norm(q).to(torch.bfloat16), _l2norm(k).to(torch.bfloat16) v = torch.randn( - batch, total, NUM_V_HEADS, HEAD_V_DIM, generator=gen, device=DEV, dtype=torch.bfloat16 + 1, total, NUM_V_HEADS, HEAD_V_DIM, generator=gen, device=DEV, dtype=torch.bfloat16 ) a_log = torch.log(torch.rand(NUM_V_HEADS, generator=gen, device=DEV) * 15 + 1) dt_bias = torch.randn(NUM_V_HEADS, generator=gen, device=DEV) - a = torch.randn(batch, total, NUM_V_HEADS, generator=gen, device=DEV) + a = torch.randn(1, total, NUM_V_HEADS, generator=gen, device=DEV) g = (-torch.exp(a_log) * torch.nn.functional.softplus(a + dt_bias)).float() - beta = torch.rand(batch, total, NUM_V_HEADS, generator=gen, device=DEV).sigmoid().float() + beta = torch.rand(1, total, NUM_V_HEADS, generator=gen, device=DEV).sigmoid().float() return q, k, v, g, beta, cu_seqlens @@ -127,57 +115,23 @@ def _gate_pool(pool: torch.Tensor, ref: tuple) -> None: def _build(spec, seed): - """`__call__`'s 13 arguments for one spec, plus the wide output buffer.""" - q, k, v, g, beta, cu_seqlens = _operands( - spec.get("seqlens"), - seed, - batch=spec.get("batch", 1), - num_tokens=spec.get("num_tokens"), - normalize_qk=spec.get("normalize_qk", True), - ) - num_seqs = q.shape[0] if cu_seqlens is None else len(cu_seqlens) - 1 - pool_dtype = spec.get("pool_dtype", torch.float32) - initial, indices, inplace = None, None, False - if spec["initial"] == "indexed": - initial = _pool(spec["slots"], seed=seed + 1, dtype=pool_dtype) - 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, dtype=pool_dtype) - wide, output = None, None - shape_o = (q.shape[0], q.shape[1], NUM_V_HEADS, HEAD_V_DIM) + """`__call__`'s 13 arguments for one spec.""" + q, k, v, g, beta, cu_seqlens = _operands(spec["seqlens"], seed) + pool = _pool(ENGINE_SLOTS, seed=seed + 1, dtype=spec.get("pool_dtype", torch.float32)) + indices = torch.tensor(spec["indices"], dtype=torch.int32, device=DEV) + output = None if spec.get("output") == "buffer": + shape_o = (q.shape[0], q.shape[1], NUM_V_HEADS, HEAD_V_DIM) output = torch.empty(shape_o, device=DEV, dtype=torch.bfloat16) - elif spec.get("output") == "window": - wide = torch.zeros(1, 200, NUM_V_HEADS, HEAD_V_DIM, device=DEV, dtype=torch.bfloat16) - output = wide[:, 50 : 50 + q.shape[1]] - args = [ - q, - k, - v, - g, - beta, - spec.get("scale"), - initial, - indices, - inplace, - spec.get("output_final_state", False), - cu_seqlens, - spec.get("l2norm", False), - output, - ] - return args, wide + return [q, k, v, g, beta, None, pool, indices, True, False, cu_seqlens, False, output] @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name - spec = cell.spec certifying(op) - args, wide = _build(spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) pool = args[6] - pool_before = None if pool is None else pool.clone() + pool_before = pool.clone() ref = op.reference(*args) o, final_state = op(*args) @@ -185,48 +139,10 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res _compare((o, final_state, pool), ref) if args[12] is not None: assert o.data_ptr() == args[12].data_ptr(), "output= must be written in place" - if wide is not None: - assert torch.count_nonzero(wide[:, :50]) == 0 and torch.count_nonzero(wide[:, 146:]) == 0 - if spec["initial"] == "positional": - assert torch.equal(pool, pool_before), "a positional initial state is read-only" - if spec["initial"] == "indexed": - named = args[7].tolist() - unnamed = [s for s in range(pool.shape[0]) if s not in set(named)] - assert torch.equal(pool[unnamed], pool_before[unnamed]), "unnamed slots written" - assert not torch.equal(pool[named], pool_before[named]), "named slots did not move" - - -def test_op_rejects_these_domains_loudly() -> None: - """Why `is_valid` guards only `output`: every other refusal is the op's.""" - q, k, v, g, beta, cu = _operands((80,), seed=67) - idx1 = torch.zeros(1, dtype=torch.int32, device=DEV) - cases = [ - (AssertionError, "does not support float32", dict(q=q.float(), k=k.float(), v=v.float())), - (AssertionError, "", dict(k=k.half())), - (AssertionError, "beta must be of shape", dict(beta=beta[0])), - (ValueError, "batch size is expected to be 1", dict(q=q.expand(2, -1, -1, -1))), - ( - ValueError, - "number of initial-state indices", - dict( - initial_state=_pool(4, seed=671), - initial_state_indices=torch.tensor([0, 1], dtype=torch.int32, device=DEV), - inplace_indexed_state_update=True, - ), - ), - (ValueError, "number of initial states", dict(initial_state=_pool(3, seed=672))), - ( - ValueError, - "require inplace_indexed_state_update=True", - dict(initial_state=_pool(1, seed=673), initial_state_indices=idx1), - ), - ] - for exc, match, override in cases: - kwargs = dict(q=q, k=k, v=v, g=g, beta=beta, cu_seqlens=cu) - kwargs.update(override) - with pytest.raises(exc, match=match): - with validating(op): - op(**kwargs) + named = args[7].tolist() + unnamed = [s for s in range(pool.shape[0]) if s not in set(named)] + assert torch.equal(pool[unnamed], pool_before[unnamed]), "unnamed slots written" + assert not torch.equal(pool[named], pool_before[named]), "named slots did not move" def test_guard_refuses_a_noncontiguous_output() -> None: @@ -246,43 +162,6 @@ def test_guard_refuses_a_noncontiguous_output() -> None: op(q, k, v, g, beta, cu_seqlens=cu, output=head_slice) -def test_guard_refuses_an_output_of_another_dtype() -> None: - """Unguarded, an fp32 buffer is filled but the returned bf16 tensor no - longer aliases it -- the allocation the buffer existed to avoid.""" - q, k, v, g, beta, cu = _operands((96,), seed=62) - buf = torch.zeros(1, 96, NUM_V_HEADS, HEAD_V_DIM, device=DEV, dtype=torch.float32) - o, _ = op(q, k, v, g, beta, cu_seqlens=cu, output=buf) - assert o.dtype == torch.bfloat16 and o.data_ptr() != buf.data_ptr(), "hazard gone" - with pytest.raises(AssertionError, match="output must have q's dtype"): - with validating(op): - op(q, k, v, g, beta, cu_seqlens=cu, output=buf) - - -def test_gate_discriminates_wrong_variants() -> None: - """Seven references wrong in one way each must all fail the `o` gate the - correct one passes. Positional path: no warm-up needed.""" - q, k, v, g, beta, cu = _operands((128, 64), seed=71) - initial = _pool(2, seed=711) - base = dict(q=q, k=k, v=v, g=g, beta=beta, scale=SCALE, initial_state=initial, cu_seqlens=cu) - o, final_state = op(**base, output_final_state=True) - ref = op.reference(**base, output_final_state=True) - _compare((o, final_state, initial), ref) - - variants = { - "no gating": dict(g=torch.zeros_like(g)), - "beta forced to 1": dict(beta=torch.ones_like(beta)), - "q and k swapped": dict(q=k, k=q), - "initial state dropped": dict(initial_state=None), - "scale left at 1": dict(scale=1.0), - "key heads reversed": dict(q=q.flip(2).contiguous(), k=k.flip(2).contiguous()), - "initial state [K, V]": dict(initial_state=initial.transpose(-1, -2).contiguous()), - } - for name, override in variants.items(): - wrong = op.reference(**{**base, **override}) - with pytest.raises(AssertionError): - _gate_output(wrong[0], ref) - - # --------------------------------------------------------------------------- # The cold call. The cells above have already benchmarked this autotune key by # the time pytest reaches it, so it runs in a fresh interpreter. diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_cu_seqlens_from_seq_lens.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_cu_seqlens_from_seq_lens.py index 27a365eaf86c..64926845e39c 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_cu_seqlens_from_seq_lens.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_cu_seqlens_from_seq_lens.py @@ -2,9 +2,9 @@ # SPDX-License-Identifier: Apache-2.0 """GPU test for the cu_seqlens_from_seq_lens entry. -Integer arithmetic, gated bitwise. Beyond the cells: a near-miss vector is -monotonic and plausible, so the tests that matter are the ones proving the -vector actually splits a packed buffer and differs from every near miss. +Integer arithmetic, gated bitwise. A near-miss vector is monotonic and +plausible, so each cell also proves the vector actually splits a packed +buffer. """ import zlib @@ -12,9 +12,6 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.cu_seqlens_from_seq_lens import ( # noqa: E501 - BATCHES, -) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.cu_seqlens_from_seq_lens import ( # noqa: E501 cu_seqlens_from_seq_lens as op, ) @@ -75,57 +72,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res assert torch.equal(packed[bounds[s] : bounds[s + 1]], piece), f"sequence {s} mis-sliced" -def test_wrong_variants_all_differ() -> None: - """Every near-miss spelling is monotonic and plausible; none is this one.""" - lengths = BATCHES["mixed_p2"] - seq_lens = _seq_lens(lengths) - out = op(seq_lens, torch.int32) - ref = op.reference(seq_lens, torch.int32) - _compare(out, ref) - - inclusive = torch.cumsum(seq_lens, dim=0, dtype=torch.int32) - no_total = torch.cat([torch.zeros(1, dtype=torch.int32, device=DEV), inclusive[:-1]]) - right_padded = torch.cat([inclusive, inclusive[-1:]]) - # The one-row-per-generation form `Mamba2Metadata.prepare` builds: a real - # cumsum over the context prefix, then a constant plus arange. - num_ctx, num_gen = 2, len(lengths) - 2 - engine = torch.zeros(len(lengths) + 1, dtype=torch.int32, device=DEV) - torch.cumsum(seq_lens[:num_ctx], dim=0, dtype=torch.int32, out=engine[1 : num_ctx + 1]) - engine[num_ctx + 1 :] = engine[num_ctx] + torch.arange( - 1, num_gen + 1, dtype=torch.int32, device=DEV - ) - for name, wrong in { - "inclusive scan": inclusive, - "total dropped": no_total, - "right-padded": right_padded, - "one row per generation": engine, - "shifted": torch.roll(out, 1), - }.items(): - with pytest.raises(AssertionError): - _compare(wrong, ref) - assert wrong.shape != ref.shape or not torch.equal(wrong, ref), name - - # The metadata form agrees exactly when every generation request emits one - # row, which is why the mistake is invisible until speculation is on. - unspec = _seq_lens(list(lengths[:num_ctx]) + [1] * num_gen) - assert torch.equal(op(unspec, torch.int32), engine) - - -def test_result_is_fresh_and_deterministic() -> None: - seq_lens = _seq_lens(BATCHES["decode_p2"]) - results = [op(seq_lens, torch.int32) for _ in range(5)] - assert len({r.data_ptr() for r in results} | {seq_lens.data_ptr()}) == 6 - for other in results[1:]: - assert torch.equal(results[0], other) - - -def test_accumulation_happens_in_the_requested_dtype() -> None: - """The hard-limit row: int64 widens an int32 input; int32 wraps silently.""" - seq_lens = _seq_lens([2_000_000_000, 2_000_000_000], torch.int32) - assert op(seq_lens, torch.int64).tolist() == [0, 2_000_000_000, 4_000_000_000] - assert op(seq_lens, torch.int32).tolist() == [0, 2_000_000_000, -294_967_296] - - def test_cuda_graph_capture_and_replay() -> None: """A captured graph follows lengths rewritten in place: only the shape reaches Python, the values are read on device and never synchronized on. @@ -150,32 +96,10 @@ def test_cuda_graph_capture_and_replay() -> None: assert torch.equal(static_out, op.reference(static_lens, torch.int32)) -def test_guard_refuses_a_seq_lens_that_is_not_1d() -> None: - """Unguarded, a [1, 3] row vector returns a length-2 result, not length 4.""" - row = _seq_lens([7, 1, 13]).unsqueeze(0) - with pytest.warns(UserWarning): - unguarded = op(row, torch.int32) - assert unguarded.shape == (2,), "the hazard this guard covers is gone" - for bad in (row, torch.zeros(3, 2, dtype=torch.int32, device=DEV), _seq_lens([3])[0]): - with pytest.raises(AssertionError, match="1-D per-request row count"): - with validating(op): - op(bad, torch.int32) - - def test_guard_refuses_non_integer_seq_lens() -> None: """Unguarded, a float32 seq_lens is truncated toward zero before the scan.""" floats = torch.tensor([7.9, 1.2, 13.6], dtype=torch.float32, device=DEV) assert op(floats, torch.int32).tolist() == [0, 7, 8, 21], "hazard gone" - for dtype in (torch.float32, torch.bfloat16, torch.bool): - with pytest.raises(AssertionError, match="seq_lens must have an integer dtype"): - with validating(op): - op(torch.ones(3, dtype=dtype, device=DEV), torch.int32) - - -def test_guard_refuses_a_non_integer_result_dtype() -> None: - """Unguarded, a float `dtype` returns a float vector the consumers misread.""" - assert op(_seq_lens([1, 2, 3]), torch.float32).dtype is torch.float32 - for dtype in (torch.float32, torch.float64, torch.bool): - with pytest.raises(AssertionError, match="dtype must be an integer dtype"): - with validating(op): - op(_seq_lens([1, 2, 3]), dtype) + with pytest.raises(AssertionError, match="seq_lens must have an integer dtype"): + with validating(op): + op(floats, torch.int32) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_dilated_short_conv_paged.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_dilated_short_conv_paged.py index 6dbc9a0eafdc..8908d558b5ad 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_dilated_short_conv_paged.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_dilated_short_conv_paged.py @@ -1,17 +1,11 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the dilated_short_conv_paged entry. - -The one engine op used here (`test_dilation_one_agrees_with_the_engine`) is an -independent oracle, not a reference: at dilation 1 this operator is exactly -`torch.ops.trtllm.causal_conv1d_fwd`, written from the other end. -""" +"""GPU test for the dilated_short_conv_paged entry.""" import zlib import pytest import torch -import torch.nn.functional as F from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.dilated_short_conv_paged import ( # noqa: E501 ENGINE_SLOTS, @@ -51,41 +45,32 @@ def _compare( def _build(spec, seed): - """`__call__`'s positional arguments for one spec, freshly allocated. + """`__call__`'s positional arguments for one spec, freshly allocated, in the + call site's form: `x` a bare `.t()` of a token-major activation, int64 + `query_start_loc`, int32 slots. The checkpoint's taps are small (its four positions average 1e-4 .. 0.0185), so the weight is scaled down rather than left at unit variance. """ g = torch.Generator(device="cuda").manual_seed(seed) channels, lengths = spec["channels"], spec["lengths"] - kernel_size, dilation = spec.get("kernel_size", KERNEL_SIZE), spec.get("dilation", DILATION) - dtype = spec.get("dtype", torch.bfloat16) - num_seqs, state_len = len(lengths), (kernel_size - 1) * dilation + num_seqs = len(lengths) def randn(*shape, scale=1.0): - return (scale * torch.randn(*shape, generator=g, device="cuda")).to(dtype) + return (scale * torch.randn(*shape, generator=g, device="cuda")).to(torch.bfloat16) starts = [0] for length in lengths: starts.append(starts[-1] + length) - qsl = torch.tensor(starts, dtype=spec.get("qsl_dtype", torch.int32), device="cuda") - if spec.get("x_form", "channel_major") == "token_major": - x = randn(starts[-1], channels).t() - else: - x = randn(channels, starts[-1]) - weight = randn(channels, 1, kernel_size, scale=0.3) - bias = randn(channels, scale=0.1) if spec.get("bias", False) else None - pool = randn(spec.get("pool_slots", 8), channels, state_len) - if spec.get("no_indices", False): - cache_indices = None - else: - slots = spec.get("slots") or tuple(range(num_seqs)) - index_dtype = spec.get("index_dtype", torch.int32) - cache_indices = torch.tensor(slots, dtype=index_dtype, device="cuda") + qsl = torch.tensor(starts, dtype=torch.int64, device="cuda") + x = randn(starts[-1], channels).t() + weight = randn(channels, 1, KERNEL_SIZE, scale=0.3) + pool = randn(spec.get("pool_slots", 8), channels, STATE_LEN) + slots = spec.get("slots") or tuple(range(num_seqs)) + cache_indices = torch.tensor(slots, dtype=torch.int32, device="cuda") initial = spec.get("initial", (True,) * num_seqs) - his = None if initial == "none" else torch.tensor(initial, dtype=torch.bool, device="cuda") - silu, pad = spec.get("silu", True), spec.get("pad_slot_id", PAD_SLOT_ID) - return [x, weight, bias, pool, qsl, cache_indices, his, dilation, silu, pad] + his = torch.tensor(initial, dtype=torch.bool, device="cuda") + return [x, weight, None, pool, qsl, cache_indices, his, DILATION, True, PAD_SLOT_ID] def _run(args): @@ -114,225 +99,6 @@ 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 at the MTP-7 verify's 32 requests of 8 tokens, replayed after the batch tensors' values change in place.""" @@ -369,23 +135,13 @@ def test_cuda_graph_capture_and_replay() -> None: assert torch.equal(live, ref_pool) -@pytest.mark.parametrize( - "case,fragment", - [ - # One per torch call that refuses: the pool's index copy, the convolution. - ("pool_dtype", "index_copy_(): self and source expected to have the same dtype"), - ("weight_2d", "weight should have at least three dimensions"), - ], -) -def test_mismatched_operands_raise(case: str, fragment: str) -> None: +def test_mismatched_operands_raise() -> None: """Refused loudly by the torch calls the entry is composed from -- why it - has no `is_valid`.""" + has no `is_valid`. Probed at a pool of another dtype (the index copy).""" args = _small((4, 4), None, seed=27) - if case == "pool_dtype": - args[3] = args[3].float() - elif case == "weight_2d": - args[1] = args[1].squeeze(1).contiguous() + args[3] = args[3].float() with pytest.raises(RuntimeError) as excinfo: with validating(op): op(*args) + fragment = "index_copy_(): self and source expected to have the same dtype" assert fragment in str(excinfo.value), str(excinfo.value) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_extract_transpose_prefill_slice.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_extract_transpose_prefill_slice.py index cf25a02c5dcf..d0e4fff2e5ad 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_extract_transpose_prefill_slice.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_extract_transpose_prefill_slice.py @@ -55,37 +55,11 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res _compare(out, op.reference(src, n, start_col, width)) -def test_the_callable_refuses_a_non_unit_column_stride() -> None: - """Loud on its own, which is why `is_valid` does not restate it.""" - src = torch.randn(256, 256, dtype=torch.bfloat16, device="cuda") - with pytest.raises(AssertionError): - op(src.t(), 128, 0, 128) - with pytest.raises(AssertionError): - op(src[:, ::2], 128, 0, 64) - - -@pytest.mark.parametrize( - "n,start_col,width,shape", - [ - (128, 0, 256, (64, 256)), # 128 rows out of 64 - (64, 200, 128, (64, 256)), # columns [200, 328) out of 256 - ], - ids=["rows_past_src", "cols_past_src"], -) -def test_guard_refuses_a_range_outside_src(n, start_col, width, shape) -> None: +def test_guard_refuses_a_range_outside_src() -> None: """The kernel masks only against the requested extents and reads past - `src` rather than failing; the guard bounds both axes.""" - src = torch.randn(*shape, dtype=torch.bfloat16, device="cuda") + `src` rather than failing; the guard bounds both axes (probed here at 128 + rows of a 64-row `src`).""" + src = torch.randn(64, 256, dtype=torch.bfloat16, device="cuda") with pytest.raises(AssertionError, match="outside src"): with validating(op): - op(src, n, start_col, width) - - -def test_guard_refuses_a_3d_src_the_callable_would_accept() -> None: - """`[rows, cols, 1]` has `stride(1) == 1`, passes the callable's assert, and - is read as a 2-D tensor through its row stride.""" - src = torch.randn(64, 256, 1, dtype=torch.bfloat16, device="cuda") - assert src.stride(1) == 1 - with pytest.raises(AssertionError, match="must be 2D"): - with validating(op): - op(src, 64, 0, 1) + op(src, 128, 0, 256) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py index 6750591cb2c1..96f88e6dbdaa 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py @@ -7,9 +7,6 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.flashinfer_chunk_gated_delta_rule import ( # noqa: E501 - QWEN38_GEOMETRY, -) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.flashinfer_chunk_gated_delta_rule import ( # noqa: E501 flashinfer_chunk_gated_delta_rule as op, ) @@ -44,16 +41,15 @@ def _compare(out: tuple, ref: tuple) -> None: def _build(spec, seed): """Keyword arguments in the domain a GDN post-conv step produces: - unit-norm q/k (unless the cell feeds raw ones), g < 0 in log space, beta a sigmoid.""" + unit-norm q/k, g < 0 in log space, beta a sigmoid.""" gen = torch.Generator(device=DEV).manual_seed(seed) num_k_heads, num_v_heads, head_k_dim, head_v_dim = spec["geometry"] seqlens = spec["seqlens"] total = sum(seqlens) q = torch.randn(1, total, num_k_heads, head_k_dim, generator=gen, device=DEV) k = torch.randn(1, total, num_k_heads, head_k_dim, generator=gen, device=DEV) - if spec.get("normalize_qk", True): - q = q / torch.sqrt(q.pow(2).sum(-1, keepdim=True) + 1e-6) - k = k / torch.sqrt(k.pow(2).sum(-1, keepdim=True) + 1e-6) + q = q / torch.sqrt(q.pow(2).sum(-1, keepdim=True) + 1e-6) + k = k / torch.sqrt(k.pow(2).sum(-1, keepdim=True) + 1e-6) q, k = q.bfloat16(), k.bfloat16() v = torch.randn( 1, total, num_v_heads, head_v_dim, generator=gen, device=DEV, dtype=torch.bfloat16 @@ -84,11 +80,11 @@ def _build(spec, seed): scale=None, initial_state=pool, initial_state_indices=indices, - inplace_indexed_state_update=spec.get("inplace", True), - output_final_state=spec.get("output_final_state", False), + inplace_indexed_state_update=True, + output_final_state=False, cu_seqlens=cu, head_first=False, - use_qk_l2norm_in_kernel=spec.get("l2norm", False), + use_qk_l2norm_in_kernel=False, output=output, state_workspace=None, ) @@ -110,70 +106,3 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res named = set(args["initial_state_indices"].tolist()) unnamed = [s for s in range(pool.shape[0]) if s not in named] assert torch.equal(pool[unnamed], pool_before[unnamed]), "unnamed slots written" - if not args["inplace_indexed_state_update"]: - assert torch.equal(pool, pool_before), "pool written without in-place update" - - -def test_repeated_calls_are_bitwise() -> None: - """Two identical calls from the same start give identical output and state.""" - results = [] - for _ in range(2): - args = _build( - dict(geometry=QWEN38_GEOMETRY, seqlens=(200, 37, 91), slots=6, indices=(5, 0, 2)), - seed=7, - ) - o, _ = op(**args) - results.append((o.clone(), args["initial_state"].clone())) - assert torch.equal(results[0][0], results[1][0]) and torch.equal(results[0][1], results[1][1]) - - -def test_matches_in_tree_triton_chunk_op() -> None: - """Same arguments through the vendored Triton chunk op this entry stands in - for: o and the written slots agree at the band. The Triton op's cold - indexed call is spent on scratch first.""" - from tensorrt_llm._torch.modules.fla.chunk import chunk_gated_delta_rule as triton_op - - def run_triton(a, pool): - return triton_op( - q=a["q"], - k=a["k"], - v=a["v"], - g=a["g"], - beta=a["beta"], - scale=None, - initial_state=pool, - initial_state_indices=a["initial_state_indices"], - inplace_indexed_state_update=True, - output_final_state=False, - cu_seqlens=a["cu_seqlens"], - head_first=False, - use_qk_l2norm_in_kernel=False, - output=None, - ) - - warm = _build(dict(geometry=QWEN38_GEOMETRY, seqlens=(64,), slots=1, indices=(0,)), seed=1) - run_triton(warm, warm["initial_state"]) - args = _build( - dict(geometry=QWEN38_GEOMETRY, seqlens=(1024, 37, 1), slots=461, indices=(460, 0, 17)), - seed=3, - ) - pool_t = args["initial_state"].clone() - o_t, _ = run_triton(args, pool_t) - o, _ = op( - args["q"], - args["k"], - args["v"], - args["g"], - args["beta"], - None, - args["initial_state"], - args["initial_state_indices"], - True, - False, - args["cu_seqlens"], - False, - False, - None, - ) - torch.testing.assert_close(o.float(), o_t.float(), rtol=RTOL, atol=ATOL_OUTPUT) - torch.testing.assert_close(args["initial_state"], pool_t, rtol=RTOL, atol=ATOL_STATE) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_gating.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_gating.py index 1a77c6d35c26..4022900963ef 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_gating.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_gating.py @@ -56,9 +56,6 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: DEV = "cuda" SOFTPLUS_THRESHOLD = 20.0 -#: log(1 + exp(x)) in fp32 is exactly 0 once exp(x) <= 2**-24. Measured: the -#: largest x with g == 0 is -16.635532, the smallest with g != 0 -16.635530. -ZERO_BELOW = -24.0 * math.log(2.0) def _params(g, num_heads=NUM_V_HEADS, dtype=torch.float32): @@ -73,22 +70,20 @@ def _build(spec, seed): """`__call__`'s positional arguments for one spec.""" g = torch.Generator(device=DEV).manual_seed(seed) rows, heads = spec["rows"], NUM_V_HEADS - 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) + return torch.randn(*shape, generator=g, device=DEV).to(torch.bfloat16) if "fill" in spec: - a = torch.full((rows, heads), spec["fill"], dtype=a_dtype, device=DEV) - elif layout == "target_rows": + a = torch.full((rows, heads), spec["fill"], dtype=torch.bfloat16, device=DEV) + elif spec.get("layout", "contiguous") == "target_rows": # `_gdn_block` splits the [T, 48] projection into prefill and decode rows. a = randn(rows + 23, heads)[23:] assert a.is_contiguous() and a.storage_offset() == 23 * heads else: a = randn(rows, heads) A_log, dt_bias = _params(g, heads) - return A_log, a, dt_bias, spec.get("beta", 1.0), spec.get("threshold", SOFTPLUS_THRESHOLD) + return A_log, a, dt_bias, 1.0, SOFTPLUS_THRESHOLD @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) @@ -100,126 +95,21 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res assert g.is_contiguous() and torch.isfinite(g).all() -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" - - -def test_op_rejects_a_non_2d_a() -> 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(2, 4, NUM_V_HEADS, device=DEV).to(torch.bfloat16) - with pytest.raises(ValueError, match="too many values to unpack"): - with validating(op): - op(A_log, a, dt_bias) - - -def test_op_rejects_cpu_tensors() -> None: - A_log, dt_bias = _params(torch.Generator(device=DEV).manual_seed(110)) - a = torch.randn(4, NUM_V_HEADS, device=DEV).to(torch.bfloat16) - with pytest.raises(ValueError, match="Pointer argument cannot be accessed from Triton"): - op(A_log.cpu(), a.cpu(), dt_bias.cpu()) - - -def _guard_case(case): - """`(args, hazard)`: the violating call and what it silently computes instead.""" +def test_guard_refuses_what_the_kernel_reads_at_the_wrong_addresses() -> None: + """Unguarded, an `a` with head stride 2 returns a finite `g` -- exactly the + one for the memory the kernel actually addressed (`row * stride(0) + + head`: the first 48 columns of each row), not the tensor it was handed; + the guard refuses it.""" g = torch.Generator(device=DEV).manual_seed(111) A_log, dt_bias = _params(g) - a = torch.randn(32, NUM_V_HEADS, generator=g, device=DEV).to(torch.bfloat16) - if case == "a_head_stride_2": - wide = torch.randn(32, BA_WIDTH, generator=g, device=DEV).to(torch.bfloat16) - bad = (A_log, wide[:, ::2], dt_bias) - # Read as `row * stride(0) + head`: the first 48 columns of each row. - return bad, (A_log, wide[:, :NUM_V_HEADS], dt_bias) - wide = torch.randn(2 * NUM_V_HEADS, generator=g, device=DEV) - if case == "A_log_stride_2": - return (wide[::2], a, dt_bias), (wide[:NUM_V_HEADS], a, dt_bias) - # Short by 8: a view into a longer buffer, so the over-read lands in known memory. - return (wide[:40], a, dt_bias), (wide[:NUM_V_HEADS], a, dt_bias) - - -@pytest.mark.parametrize( - "case,match", - [ - ("a_head_stride_2", "head axis must be unit-stride"), - ("A_log_stride_2", "must be contiguous 1-D"), - ("A_log_short", "at least a.shape"), - ], -) -def test_guard_refuses_what_the_kernel_reads_at_the_wrong_addresses(case, match) -> None: - """Unguarded, each call returns a finite `g` -- exactly the one for the - memory the kernel actually addressed, not the tensor it was handed; the - guard refuses it.""" - bad, addressed = _guard_case(case) + wide = torch.randn(32, BA_WIDTH, generator=g, device=DEV).to(torch.bfloat16) + bad = (A_log, wide[:, ::2], dt_bias) + addressed = (A_log, wide[:, :NUM_V_HEADS], dt_bias) unguarded = op(*bad) assert torch.isfinite(unguarded).all() assert torch.equal(unguarded, op(*addressed)), "the hazard this guard covers is gone" - if "short" not in case: - with pytest.raises(AssertionError): - _compare(unguarded, op.reference(*bad)) - with pytest.raises(AssertionError, match=match): + with pytest.raises(AssertionError): + _compare(unguarded, op.reference(*bad)) + with pytest.raises(AssertionError, match="head axis must be unit-stride"): with validating(op): op(*bad) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_post_conv.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_post_conv.py index d4280185fe90..01b35ad0d378 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_post_conv.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_post_conv.py @@ -82,142 +82,62 @@ def _compare(out: tuple, ref: tuple) -> None: def _build(spec, seed): - """`__call__`'s positional arguments, then its keyword `beta_dtype`.""" + """`__call__`'s positional arguments (its keyword `beta_dtype` left at fp32).""" g = torch.Generator(device=DEV).manual_seed(seed) n_p, n_d = spec["num_prefill"], spec["num_decode"] n = n_p + n_d - layout = spec.get("layout", "contiguous") def randn(*shape): return torch.randn(*shape, generator=g, device=DEV).to(DTYPE) - if layout == "in_proj": - # `_gdn_block`: one [T, 16480] GEMM output; q|k|v, a and b are its column - # views, the decode half the row split of q|k|v, the prefill half the - # transposing extraction of the rest. - proj = randn(n, IN_PROJ_WIDTH) - mixed = proj[:, :QKV_DIM] - prefill = mixed[:n_p].t().contiguous() - decode = mixed[n_p:] if n_d else None - a = proj[:, QKV_DIM + Z_WIDTH : QKV_DIM + Z_WIDTH + NUM_V_HEADS] - b = proj[:, QKV_DIM + Z_WIDTH + NUM_V_HEADS :] - assert a.stride(0) == b.stride(0) == IN_PROJ_WIDTH - assert decode is None or decode.stride() == (IN_PROJ_WIDTH, 1) - 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) - else: - prefill = randn(QKV_DIM, n_p) - decode = randn(n_d, QKV_DIM) if n_d else None - a, b = randn(n, NUM_V_HEADS), randn(n, NUM_V_HEADS) + # `_gdn_block`: one [T, 16480] GEMM output; q|k|v, a and b are its column + # views, the decode half the row split of q|k|v, the prefill half the + # transposing extraction of the rest. + proj = randn(n, IN_PROJ_WIDTH) + mixed = proj[:, :QKV_DIM] + prefill = mixed[:n_p].t().contiguous() + decode = mixed[n_p:] if n_d else None + a = proj[:, QKV_DIM + Z_WIDTH : QKV_DIM + Z_WIDTH + NUM_V_HEADS] + b = proj[:, QKV_DIM + Z_WIDTH + NUM_V_HEADS :] + assert a.stride(0) == b.stride(0) == IN_PROJ_WIDTH + assert decode is None or decode.stride() == (IN_PROJ_WIDTH, 1) - gate_dtype = spec.get("gate_dtype", torch.float32) lo, hi = A_LOG_RANGE - A_log = (torch.rand(NUM_V_HEADS, generator=g, device=DEV) * (hi - lo) + lo).to(gate_dtype) + A_log = torch.rand(NUM_V_HEADS, generator=g, device=DEV) * (hi - lo) + lo lo, hi = DT_BIAS_RANGE - dt_bias = (torch.rand(NUM_V_HEADS, generator=g, device=DEV) * (hi - lo) + lo).to(gate_dtype) + dt_bias = torch.rand(NUM_V_HEADS, generator=g, device=DEV) * (hi - lo) + lo if "dt_bias_fill" in spec: - dt_bias = torch.full((NUM_V_HEADS,), spec["dt_bias_fill"], device=DEV).to(gate_dtype) - args = [prefill, decode, a, b, A_log, dt_bias, *HEADS, L2_NORM_EPS, SOFTPLUS_THRESHOLD] - return args, spec.get("beta_dtype", torch.float32) + dt_bias = torch.full((NUM_V_HEADS,), spec["dt_bias_fill"], device=DEV) + return [prefill, decode, a, b, A_log, dt_bias, *HEADS, L2_NORM_EPS, SOFTPLUS_THRESHOLD] @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name certifying(op) - args, beta_dtype = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) if "dt_bias_fill" in cell.spec: # Every head strictly on the side of the threshold the cell names. x = args[2].float() + args[5].float() above = cell.spec["dt_bias_fill"] > SOFTPLUS_THRESHOLD assert bool((x > SOFTPLUS_THRESHOLD).all() if above else (x <= SOFTPLUS_THRESHOLD).all()) before = [t.clone() if isinstance(t, torch.Tensor) else t for t in args] - out = op(*args, beta_dtype=beta_dtype) + out = op(*args) for arg, old in zip(args, before): if isinstance(arg, torch.Tensor): assert torch.equal(arg, old), "inputs are read-only" - _compare(out, op.reference(*args, beta_dtype=beta_dtype)) - - -def test_g_gate_discriminates_wrong_variants() -> None: - """`g`'s floor is derived, so prove it still separates. Each variant is a - plausible way to get the gate wrong, and must land >1e4x outside it.""" - args, _ = _build(dict(num_prefill=96, num_decode=0), seed=106) - prefill, decode, a, b, A_log, dt_bias = args[:6] - out = op(*args) - ref = op.reference(*args) - _compare(out, ref) - g_ref = ref[3].double() - atol = 4 * 2.0**-22 * SOFTPLUS_THRESHOLD * math.exp(A_LOG_RANGE[1]) - observed = (out[3].double() - g_ref).abs().max().item() - assert observed < atol, f"correct result {observed} outside its own gate {atol}" - - x = a.double() + dt_bias.double() - softplus = torch.where(x <= SOFTPLUS_THRESHOLD, torch.log1p(torch.exp(x)), x) - - def g_of(*overrides): - return op.reference(*[*overrides, *args[len(overrides) :]])[3].double() - - variants = { - "sign flipped": -g_ref, - "dt_bias omitted": g_of(prefill, decode, a, b, A_log, torch.zeros_like(dt_bias)), - "A_log not exponentiated": (-A_log.double() * softplus).unsqueeze(0), - "softplus replaced by identity": (-torch.exp(A_log.double()) * x).unsqueeze(0), - "a and b swapped": g_of(prefill, decode, b, a), - } - for name, variant in variants.items(): - distance = (variant - g_ref).abs().max().item() - assert distance > 1e4 * atol, f"{name} lands {distance / atol:.1f}x the gate" - with pytest.raises(AssertionError): - _compare((*out[:3], variant.float(), out[4]), ref) - - -def test_qk_normalization_is_not_a_pass_through() -> None: - """The default bf16 band on q/k separates a missing norm and a q/k swap.""" - args, _ = _build(dict(num_prefill=96, num_decode=0), seed=107) - out = op(*args) - ref = op.reference(*args) - _compare(out, ref) - q_raw = args[0].t()[:, : NUM_K_HEADS * HEAD_K_DIM].reshape(1, -1, NUM_K_HEADS, HEAD_K_DIM) - for wrong in ((q_raw.contiguous(), *out[1:]), (out[1], *out[1:])): - with pytest.raises(AssertionError): - _compare(wrong, ref) + _compare(out, op.reference(*args)) def test_guard_refuses_strided_gate_parameters() -> None: - """Unguarded, a stride-2 view is read as `ptr + head` -- the interleaved - zeros -- and returns a finite wrong g.""" - args, _ = _build(dict(num_prefill=32, num_decode=0), seed=108) - for position in (4, 5): - strided = torch.stack([args[position], torch.zeros_like(args[position])], dim=1)[:, 0] - assert strided.stride() == (2,) - bad = list(args) - bad[position] = strided - with pytest.raises(AssertionError): - _compare(op(*bad), op.reference(*bad)) - with pytest.raises(AssertionError, match="must be contiguous 1-D"): - with validating(op): - op(*bad) - - -def test_guard_refuses_short_gate_parameters_and_a_b() -> None: - """Each of these is an unmasked out-of-bounds read, so the hazard is not - replayed here; the contract's measurements are in the entry's note.""" - args, _ = _build(dict(num_prefill=32, num_decode=8), seed=109) - for position, match in ((4, "at least num_v_heads"), (5, "at least num_v_heads")): - bad = list(args) - bad[position] = args[position][: NUM_V_HEADS - 8].clone() - with pytest.raises(AssertionError, match=match): - with validating(op): - op(*bad) - for position in (2, 3): - bad = list(args) - # 32 rows cover the prefill block but not prefill + decode = 40. - bad[position] = args[position][:32].clone() - with pytest.raises(AssertionError, match="at least num_tokens=40"): - with validating(op): - op(*bad) + """Unguarded, a stride-2 `A_log` view is read as `ptr + head` -- the + interleaved zeros -- and returns a finite wrong g.""" + args = _build(dict(num_prefill=32, num_decode=0), seed=108) + strided = torch.stack([args[4], torch.zeros_like(args[4])], dim=1)[:, 0] + assert strided.stride() == (2,) + bad = list(args) + bad[4] = strided + with pytest.raises(AssertionError): + _compare(op(*bad), op.reference(*bad)) + with pytest.raises(AssertionError, match="must be contiguous 1-D"): + with validating(op): + op(*bad) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py index 341f11ff8c21..99a60e6e074d 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py @@ -18,8 +18,7 @@ advanced by `p += accepted` alone. The twin's state is the plain token-by-token recurrence over every accepted token, independent of the wrap, the double buffer and the commit, so a wrong commit or wrong - bookkeeping surfaces as an output that drifts off it -- - `test_the_twin_sees_a_broken_protocol` shows it does. + bookkeeping surfaces as an output that drifts off it. """ from __future__ import annotations @@ -29,9 +28,7 @@ import pytest import torch import torch.nn.functional as F -from triton.compiler.errors import CompilationError -import tensorrt_llm._torch.modules.fla.cached_replay 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_cached_replay_update import ( # noqa: E501 ALL_LAYER_COMMIT_MIN_BATCH, HEAD_K_DIM, @@ -65,7 +62,6 @@ H, HV, K, V = NUM_K_HEADS, NUM_V_HEADS, HEAD_K_DIM, HEAD_V_DIM W = PACKED_QKV_WIDTH HIST = HISTORY_SIZE -OUT_SENTINEL = -7.0 BF16_ULP = ULP[torch.bfloat16] #: This checkpoint's 36 GDN layers of `linear_attn.A_log` / `linear_attn.dt_bias` @@ -198,9 +194,7 @@ def __init__(self, spec, seed, exact=True): else: chosen = torch.randperm(S, generator=self.cpu)[:n] chosen = chosen.to(torch.int32) - for row in spec.get("pad", ()): - chosen[row] = -1 - live_rows = chosen >= 0 + live_rows = torch.ones(n, dtype=torch.bool) dup = tuple(spec.get("dup", ())) if dup: chosen[list(dup)] = int(chosen[dup[0]]) @@ -270,22 +264,9 @@ def inputs(self): q = rows[..., : H * K].view(n, t, H, K) k = rows[..., H * K : 2 * H * K].view(n, t, H, K) v = rows[..., 2 * H * K :].view(n, t, HV, V) - if self.spec.get("packed", True) is False: - packed = None # the mixed batch's form: the views alone, copied by the op - if self.spec.get("ab", "split") == "strided": - ab = self._randn(n * t, 2 * HV, dtype=torch.bfloat16) - b, a = ab[:, :HV].view(n, t, HV), ab[:, HV:].view(n, t, HV) - else: - a = self._randn(n, t, HV, dtype=torch.bfloat16) - b = self._randn(n, t, HV, dtype=torch.bfloat16) - out_buf = output = None - lo = 0 if self.spec.get("output") == "prefix" else 5 - if self.spec.get("output") in ("prefix", "suffix"): - # A slice of a longer token-major buffer: its first rows ("prefix"), - # or the rows after 5 context tokens, as a mixed batch hands it over. - out_buf = torch.full((n * t + 5, HV, V), OUT_SENTINEL, dtype=torch.bfloat16, device=DEV) - output = out_buf[lo : lo + n * t].view(n, t, HV, V) - return dict(q=q, k=k, v=v, a=a, b=b, packed=packed, output=output, out_buf=out_buf, lo=lo) + a = self._randn(n, t, HV, dtype=torch.bfloat16) + b = self._randn(n, t, HV, dtype=torch.bfloat16) + return dict(q=q, k=k, v=v, a=a, b=b, packed=packed) def kwargs(self, layer, act): return dict( @@ -307,12 +288,12 @@ def kwargs(self, layer, act): use_qk_l2norm_in_kernel=True, A_log=self.a_log[layer], dt_bias=self.dt_bias[layer], - launch_with_pdl=self.spec.get("pdl", False), + launch_with_pdl=False, replay_work_items=self.work_items, n_writes=self.n_writes, packed_qkv=act["packed"], use_all_layer_commit=self.path == "all_layer", - output=act["output"], + output=None, ) def exact_kwargs(self, layer, kw): @@ -368,8 +349,6 @@ def draw_acceptance(self): rows = self.live.numel() if mode == "full": return torch.full((rows,), self.t, dtype=torch.int32, device=DEV) - if mode == "one": - return torch.ones(rows, dtype=torch.int32, device=DEV) assert mode == "random", mode return torch.randint(1, self.t + 1, (rows,), generator=self.cpu, dtype=torch.int32).to(DEV) @@ -392,19 +371,6 @@ def writes(self): return mask, s[wrap], int((p + self.t == HIST).sum()) -def _deviation(got, want): - """The largest element deviation relative to its block's absmax (the - quantity `_compare` gates).""" - block = got.shape[-2] * got.shape[-1] - g, w = got.float().reshape(-1, block), want.float().reshape(-1, block) - keep = ~torch.isnan(w).any(1) - g, w = g[keep], w[keep] - if not w.numel(): - return 0.0 - scale = w.abs().amax(1, keepdim=True).clamp_min(1e-30) - return ((g - w).abs() / scale).max().item() - - def _check_duplicate_padding_slot(cell) -> None: """The last rows of a batch share one slot, as CUDA-graph padding repeats the manager's dummy slot: the rows on unique slots, and every slot but the shared one, are held to @@ -449,7 +415,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res if "dup" in cell.spec: return _check_duplicate_padding_slot(cell) bench = _Bench(cell.spec, seed=_seed(cell.why)) - pad = list(cell.spec.get("pad", ())) wraps = appends = to_the_end = 0 for _ in range(bench.steps): mask, wrapped, at_end = bench.writes() @@ -468,7 +433,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res k=k.clone(), G=G.clone(), old_beta=bench.old_dA[layer].clone(), - out_buf=None if act["out_buf"] is None else act["out_buf"].clone(), ) ref = op.reference(**kw) exact = op.reference(**bench.exact_kwargs(layer, kw)) @@ -492,15 +456,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res assert torch.equal(pool[untouched], before["pool"][untouched]), ( "the pool moved where the call commits nothing (a partitioned call commits none)" ) - if kw["output"] is not None: - assert got.data_ptr() == kw["output"].data_ptr(), "output must be written in place" - outside = torch.ones(act["out_buf"].shape[0], dtype=torch.bool, device=DEV) - outside[act["lo"] : act["lo"] + bench.n * bench.t] = False - assert torch.equal(act["out_buf"][outside], before["out_buf"][outside]), ( - "rows of the output buffer outside the view were touched" - ) - for row in pad: - assert (got[row] == OUT_SENTINEL).all(), f"padding row {row} written" bench.adopt_exact(layer, exact) after_commit.append((got, exact, committed)) bench.commit() @@ -525,23 +480,14 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res # ── CUDA graphs ─────────────────────────────────────────────────────────── -#: The in-kernel commit at one sequence and the smallest token tile, the all-layer commit at -#: the envelope's 32 requests and the largest, and a bf16 checkpoint at the partition threshold. -_GRAPH_CASES = [(1, 2, torch.float32), (32, 8, torch.float32), (16, 4, torch.bfloat16)] - - -@pytest.mark.parametrize( - "n,t,state_dtype", - _GRAPH_CASES, - ids=[f"{n}-{t}-{'fp32' if d is torch.float32 else 'bf16'}" for n, t, d in _GRAPH_CASES], -) -def test_cuda_graph_replay_on_new_data_is_bitwise_eager(n, t, state_dtype) -> None: - """Two layers captured with, from 16 sequences on, the all-layer commit after them; - replayed three times on new activations, with acceptance in between, against the same - steps run eagerly on a twin state.""" +def test_cuda_graph_replay_on_new_data_is_bitwise_eager() -> None: + """Two layers of 32 sequences x 8 tokens over an fp32 checkpoint captured with the + all-layer commit after them; replayed three times on new activations, with acceptance + in between, against the same steps run eagerly on a twin state.""" + n, t = 32, 8 assert ALL_LAYER_COMMIT_MIN_BATCH == CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE - path = "all_layer" if n >= ALL_LAYER_COMMIT_MIN_BATCH else "small" - spec = dict(path=path, n=n, t=t, layers=2, slots=40, state_dtype=state_dtype) + assert n >= ALL_LAYER_COMMIT_MIN_BATCH + spec = dict(path="all_layer", n=n, t=t, layers=2, slots=40, state_dtype=torch.float32) seed = 7000 + 10 * n + t eager, graph, warm = (_Bench(spec, seed, exact=False) for _ in range(3)) acts = [warm.inputs() for _ in range(2)] @@ -580,178 +526,6 @@ def test_cuda_graph_replay_on_new_data_is_bitwise_eager(n, t, state_dtype) -> No graph.accept(accepted) -# ── the harness can see what it exists to see ───────────────────────────── - - -def _drive(bench, fault): - """Run `bench` with one protocol fault injected; return the first step at - which the twin comparison rejected an output, or None.""" - for step in range(bench.steps): - bench.prepare() - outs = [] - for layer in range(bench.layers): - kw = bench.kwargs(layer, bench.inputs()) - if fault == "committed twice": - kw = dict(kw, replay_work_items=None, n_writes=None, use_all_layer_commit=False) - exact = op.reference(**bench.exact_kwargs(layer, kw)) - bench.adopt_exact(layer, exact) - outs.append((op(**kw), exact[0])) - if fault != "commit skipped": - bench.commit() - try: - for got, want in outs: - _compare((got, None, None, None, None), (want, None, None, None, None)) - except AssertionError: - return step - accepted = bench.draw_acceptance() - if fault == "rejected tokens kept": - _advance_replay_state(bench.meta, bench.live, torch.full_like(accepted, bench.t)) - bench.x_pnat[bench.live.long()] += accepted - else: - bench.accept(accepted) - return None - - -@pytest.mark.parametrize( - "fault", ("none", "commit skipped", "committed twice", "rejected tokens kept") -) -def test_the_twin_sees_a_broken_protocol(fault) -> None: - """The twin comparison is what certifies the commit and the bookkeeping, so - it has to fail when either is broken: the partitioned path with the - manager's commit left out; a 16-row batch called without work items while - the manager still commits (the checkpoint advanced twice); and the small - path advanced as if every drafted token had been accepted. Without a fault - the same run passes. At T = 8, the largest token tile.""" - path = "small" if fault == "rejected tokens kept" else "all_layer" - n = 8 if path == "small" else 16 - bench = _Bench(dict(path=path, n=n, t=8, steps=4), seed=93) - first = _drive(bench, fault) - if fault == "none": - assert first is None - else: - assert first is not None, f"{fault} passed the twin comparison" - - -@pytest.mark.parametrize("state_dtype", (torch.float32, torch.bfloat16), ids=("fp32", "bf16")) -def test_the_gate_rejects_near_misses(state_dtype) -> None: - """`_compare` against plausible wrong answers, each the entry's own - reference driven on transformed inputs, at a call where one row wraps (and - commits) and the others append -- at each checkpoint dtype's own state gate.""" - bench = _Bench(dict(path="small", n=8, t=4, state_dtype=state_dtype), seed=91, exact=False) - kw = bench.kwargs(0, bench.inputs()) - pool, u, k, G = (x.clone() for x in bench.state(0)) - base = dict(kw, ssm_states=pool, old_u=u, old_k=k, old_G=G) - ref = op.reference(**base) - got = op(**kw) - out = (got, *bench.state(0)) - _compare(out, ref) - p, b = kw["prev_num_accepted_tokens"], kw["cache_buf_idx"] - - not_restarted = ref[4].clone() - _, wrapped, _ = bench.writes() - for s in wrapped.tolist(): - g_start = G[s, b[s].long(), :, p[s].long() - 1] - not_restarted[s, 1 - b[s].long(), :, : bench.t] += g_start[:, None] - rivals = { - "history ignored": op.reference(**dict(base, prev_num_accepted_tokens=torch.zeros_like(p))), - "the other history half": op.reference(**dict(base, cache_buf_idx=1 - b)), - "one pending token dropped": op.reference( - **dict(base, prev_num_accepted_tokens=(p - 1).clamp(min=0)) - ), - "checkpoint zeroed": op.reference(**dict(base, ssm_states=torch.zeros_like(pool))), - "G not restarted at a wrap": ref[:4] + (not_restarted,), - } - for rival in rivals.values(): - with pytest.raises(AssertionError): - _compare(out, rival) - - -# ── properties beyond one reference comparison ──────────────────────────── - - -def test_the_wrapped_kernels_carry_no_autotuner() -> None: - """Plain `triton.jit`: no trial launch can consume the in-place state on a - process's first call.""" - for kernel in (_impl._cached_replay_kernel, _impl._cached_replay_layered_commit_kernel): - names = [cls.__name__ for cls in type(kernel).__mro__] - assert "Autotuner" not in names and "Heuristics" not in names, names - - -def _clone_state(bench): - return {name: getattr(bench, name).clone() for name in ("backing", "old_x", "old_B", "old_dt")} - - -def _restore_state(bench, saved): - for name, value in saved.items(): - getattr(bench, name).copy_(value) - - -@pytest.mark.parametrize("state_dtype", (torch.float32, torch.bfloat16), ids=("fp32", "bf16")) -def test_the_two_commit_paths_are_bitwise_interchangeable(state_dtype) -> None: - """The same 16 rows through the in-kernel commit, and through the - partitioned call plus the manager's all-layer commit: bit for bit, at either - checkpoint dtype.""" - bench = _Bench( - dict(path="all_layer", n=16, t=4, layers=2, state_dtype=state_dtype), seed=95, exact=False - ) - acts = [bench.inputs() for _ in range(2)] - saved = _clone_state(bench) - small = [ - op( - **dict( - bench.kwargs(i, acts[i]), - replay_work_items=None, - n_writes=None, - use_all_layer_commit=False, - ) - ) - for i in range(2) - ] - after_small = _clone_state(bench) - _restore_state(bench, saved) - bench.prepare() - assert 0 < int(bench.n_writes) < bench.n, "want both wrapping and appending rows" - partitioned = [op(**bench.kwargs(i, acts[i])) for i in range(2)] - bench.commit() - for a, b in zip(small, partitioned): - assert torch.equal(a, b) - for name, value in after_small.items(): - assert torch.equal(getattr(bench, name), value), name - - -def test_packed_qkv_supersedes_q_k_v_and_launch_knobs_are_neutral() -> None: - """With `packed_qkv` the kernel reads q/k/v from it alone -- their own - contents are never read; the separate contiguous form, PDL on or off, and a - second run are all bitwise the same call.""" - bench = _Bench(dict(path="small", n=6, t=3, layers=1), seed=97, exact=False) - act = bench.inputs() - saved = _clone_state(bench) - - def run(**override): - _restore_state(bench, saved) - out = op(**dict(bench.kwargs(0, act), **override)).clone() - return out, _clone_state(bench) - - base = run() - q, k, v = act["q"], act["k"], act["v"] - variants = { - "separate q/k/v": run( - q=q.contiguous(), k=k.contiguous(), v=v.contiguous(), packed_qkv=None - ), - "q/k/v contents zeroed": run( - q=torch.zeros_like(q.contiguous()), - k=torch.zeros_like(k.contiguous()), - v=torch.zeros_like(v.contiguous()), - ), - "launch_with_pdl=True": run(launch_with_pdl=True), - "second run": run(), - } - for name, (out, state) in variants.items(): - assert torch.equal(out, base[0]), name - for key, value in state.items(): - assert torch.equal(value, base[1][key]), (name, key) - - # ── guards: accepted and silently wrong ────────────────────────────────── @@ -771,109 +545,3 @@ def test_a_short_state_indices_reads_past_its_end_and_is_refused() -> None: with pytest.raises(AssertionError, match="state_indices has shape"): with validating(op): op(**kw) - - -def test_a_history_off_history_size_is_misread_and_is_refused() -> None: - """Unguarded, 16-deep buffers called with history_size=12: the kernel strides - them as 12 deep and reads other slots' entries. The same call on 12-deep - copies of the same history is the right answer. (12, not 8: below 9 the - kernel does not compile -- see `test_op_refuses_what_it_refuses`.)""" - bench = _Bench(dict(path="small", n=4, t=4, layers=1), seed=103, exact=False) - live = bench.live.long() - bench.pnat[live] = torch.tensor([1, 2, 3, 4], dtype=torch.int32, device=DEV) - kw = bench.kwargs(0, bench.inputs()) - shallow = dict( - kw, - old_u=kw["old_u"][:, :, :12].clone(), - old_k=kw["old_k"][:, :, :12].clone(), - old_G=kw["old_G"][..., :12].clone(), - ssm_states=kw["ssm_states"].clone(), - history_size=12, - ) - ref = op.reference(**shallow) - right = op(**shallow) - _compare((right, None, None, None, None), (ref[0], None, None, None, None)) - wrong = op(**dict(kw, history_size=12)) - gap = _deviation(wrong, right) / BF16_ULP - assert gap > 64, "the hazard this guard covers is gone" - with pytest.raises(AssertionError, match="old_u has shape"): - with validating(op): - op(**dict(kw, history_size=12)) - - -@pytest.mark.parametrize( - "name, form", - ( - ("history", "fewer slots"), - ("prev_num_accepted_tokens", "fewer slots"), - ("cache_buf_idx", "strided"), - ), -) -def test_bookkeeping_that_does_not_cover_the_pool_is_refused(name, form) -> None: - """The kernel addresses the history buffers, `prev_num_accepted_tokens` and - `cache_buf_idx` by the pool's slot id: one with fewer slots than `ssm_states` is read and - written past its end for the slots beyond it, and a strided per-slot vector is read at the - wrong elements. Only the refusal runs; the unguarded call is out of bounds.""" - bench = _Bench(dict(path="small", n=4, t=2, layers=1), seed=109, exact=False) - kw = bench.kwargs(0, bench.inputs()) - last = bench.slots - 1 - if name == "history": - kw = dict(kw, old_u=kw["old_u"][:last], old_k=kw["old_k"][:last], old_G=kw["old_G"][:last]) - match = "history buffers have" - else: - vec = kw[name] - if form == "fewer slots": - kw[name] = vec[:last] - else: - kw[name] = torch.stack([vec, vec], 1)[:, 0] - match = f"{name} has shape" - with pytest.raises(AssertionError, match=match): - with validating(op): - op(**kw) - - -def test_op_refuses_what_it_refuses() -> None: - """Loud refusals; the guard admits each so the op's own error surfaces.""" - bench = _Bench(dict(path="all_layer", n=16, t=4, layers=1), seed=105, exact=False) - bench.prepare() - kw = bench.kwargs(0, bench.inputs()) - with pytest.raises(ValueError, match="Partitioned replay requires the all-layer commit"): - with validating(op): - op(**dict(kw, use_all_layer_commit=False)) - with pytest.raises(ValueError, match="use_all_layer_commit requires replay work items"): - with validating(op): - op(**dict(kw, replay_work_items=None, n_writes=None)) - with pytest.raises(ValueError, match="must either both be set or both be None"): - with validating(op): - op(**dict(kw, n_writes=None)) - with pytest.raises(AssertionError, match="output must be contiguous"): - wide = torch.empty(16, 4, HV, 2 * V, dtype=torch.bfloat16, device=DEV) - with validating(op): - op(**dict(kw, output=wide[..., :V])) - with pytest.raises(AssertionError, match="g head dimension must be contiguous"): - with validating(op): - op(**dict(kw, g=kw["g"].transpose(1, 2).contiguous().transpose(1, 2))) - - def shallow(depth): - return dict( - kw, - history_size=depth, - old_u=kw["old_u"][:, :, :depth].contiguous(), - old_k=kw["old_k"][:, :, :depth].contiguous(), - old_G=kw["old_G"][..., :depth].contiguous(), - ) - - with pytest.raises(AssertionError): # `history_size >= T`, a bare assert - with validating(op): - op(**shallow(2)) - # The history tile is next_pow2(history_size) wide and is a `tl.dot` - # reduction, which needs 16: history_size 8 (>= T = 4) fails to compile. - with pytest.raises(CompilationError, match="K >= 16"): - with validating(op): - op(**shallow(8)) - # Nine rows per sequence (max_draft_len 8) is past the kernel's `T <= 8`, a bare assert - # before the launch. - wide = _Bench(dict(path="small", n=2, t=9, layers=1), seed=106, exact=False) - with pytest.raises(AssertionError): - with validating(op): - op(**wide.kwargs(0, wide.inputs())) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_update.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_update.py index 8da57911d100..3b2041b63919 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_update.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_update.py @@ -5,14 +5,12 @@ The axis this entry exists for is more than one token per sequence per step -- the speculative verify call -- so the cells are that call as the target makes it (per-step states into the rollback buffer, the pool untouched) at both pool -dtypes. The write-back sink is driven only by the property and guard tests below. +dtypes. The write-back sink is driven only by the cold-call test below. -Two tests answer "is this harness blind" rather than a correctness question: -the neighbouring `chunk_gated_delta_rule` is wrong on the first indexed +The neighbouring `chunk_gated_delta_rule` was once wrong on the first indexed in-place call of a process (an autotuner trial launch consumes the state), so `test_cold_call_in_a_fresh_process_is_correct` measures the first call of a -process that has made none, and `test_the_gate_rejects_a_doubly_applied_state` -shows `_compare` rejects exactly the damage that mechanism leaves. +process that has made none. """ from __future__ import annotations @@ -26,7 +24,6 @@ import pytest import torch -import tensorrt_llm._torch.modules.fla.fused_recurrent as _impl from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_recurrent_gated_delta_rule_update import ( # noqa: E501 HEAD_K_DIM, HEAD_V_DIM, @@ -74,67 +71,49 @@ def _compare(out: tuple, ref: tuple) -> None: DEV = "cuda" -REPEAT = NUM_V_HEADS // NUM_K_HEADS # 3 MIXED_QKV_WIDTH = 2 * NUM_K_HEADS * HEAD_K_DIM + NUM_V_HEADS * HEAD_V_DIM # 10240 CACHE_SENTINEL = -3.0 def _build(spec, seed): - """`__call__`'s arguments by name for one spec, plus what the locks need. + """`__call__`'s arguments by name for one spec, plus the live slots. q/k/v are views over one packed `[T, 10240]` row block (or contiguous copies, as the target passes them); `g` is built the way the gating step builds it -- negative, so `exp(g) < 1` -- and `beta` is a sigmoid. """ gen = torch.Generator(device=DEV).manual_seed(seed) - mode, n = spec["mode"], spec["n"] - lengths = spec.get("lengths") - batch, tokens = (n, spec["p"]) if lengths is None else (1, sum(lengths)) - total = batch * tokens + mode, n, tokens = spec["mode"], spec["n"], spec["p"] + total = n * tokens def randn(*shape, dtype=torch.float32): return torch.randn(*shape, generator=gen, device=DEV, dtype=dtype) mixed = randn(total, MIXED_QKV_WIDTH, dtype=torch.bfloat16) kd = NUM_K_HEADS * HEAD_K_DIM - q = mixed[:, :kd].view(batch, tokens, NUM_K_HEADS, HEAD_K_DIM) - k = mixed[:, kd : 2 * kd].view(batch, tokens, NUM_K_HEADS, HEAD_K_DIM) - v = mixed[:, 2 * kd :].view(batch, tokens, NUM_V_HEADS, HEAD_V_DIM) + q = mixed[:, :kd].view(n, tokens, NUM_K_HEADS, HEAD_K_DIM) + k = mixed[:, kd : 2 * kd].view(n, tokens, NUM_K_HEADS, HEAD_K_DIM) + v = mixed[:, 2 * kd :].view(n, tokens, NUM_V_HEADS, HEAD_V_DIM) if spec.get("qkv") == "contiguous": q, k, v = q.contiguous(), k.contiguous(), v.contiguous() a, a_log, dt_bias = randn(total, NUM_V_HEADS), randn(NUM_V_HEADS), randn(NUM_V_HEADS) g = -torch.exp(a_log) * torch.nn.functional.softplus(a + dt_bias) - g = g.view(batch, tokens, NUM_V_HEADS) + g = g.view(n, tokens, NUM_V_HEADS) b = randn(total, NUM_V_HEADS, dtype=torch.bfloat16) - beta_kind = spec.get("beta", "fp32") - beta = (b.sigmoid() if beta_kind == "bf16" else b.float().sigmoid()).view( - batch, tokens, NUM_V_HEADS - ) - if beta_kind == "channel": - beta = beta[..., None].expand(batch, tokens, NUM_V_HEADS, HEAD_V_DIM).contiguous() + beta = b.sigmoid() if spec.get("beta", "fp32") == "bf16" else b.float().sigmoid() + beta = beta.view(n, tokens, NUM_V_HEADS) pool_slots = spec.get("pool_slots", SLOTS) - pool = layered = None - if pool_slots is not None: - pool = randn(pool_slots, NUM_V_HEADS, HEAD_V_DIM, HEAD_K_DIM) * 0.05 - pool = pool.to(spec.get("pool_dtype", torch.float32)) - if spec.get("pool_view") == "layer_major": - layered = torch.zeros((3,) + pool.shape, dtype=pool.dtype, device=DEV) - layered[1].copy_(pool) - pool = layered[1] - assert pool.is_contiguous() + pool = randn(pool_slots, NUM_V_HEADS, HEAD_V_DIM, HEAD_K_DIM) * 0.05 + pool = pool.to(spec.get("pool_dtype", torch.float32)) slots = spec.get("slots") or tuple(range(n)) - indices = torch.tensor(slots, dtype=spec.get("index_dtype", torch.int32), device=DEV) - cu_seqlens = None - if lengths is not None: - bounds = torch.tensor((0,) + tuple(lengths)).cumsum(0) - cu_seqlens = bounds.to(device=DEV, dtype=spec.get("cu_dtype", torch.int64)) + indices = torch.tensor(slots, dtype=torch.int32, device=DEV) buffer = cache_steps = None - if mode in ("verify", "both"): - rows = spec.get("cache_rows", SLOTS if pool_slots is None else pool_slots) + if mode == "verify": + rows = spec.get("cache_rows", pool_slots) cache_steps = tokens buffer = torch.full( (rows, cache_steps, NUM_V_HEADS, HEAD_V_DIM, HEAD_K_DIM), @@ -143,17 +122,6 @@ def randn(*shape, dtype=torch.float32): device=DEV, ) - output = out_buf = None - if spec.get("output") == "prefix": - # A length-prefix slice of a longer row-major buffer, as the stock mixer - # passes it: contiguous, and the tail is not offered. - out_buf = torch.full( - (total + 5, NUM_V_HEADS, HEAD_V_DIM), -7.0, dtype=torch.bfloat16, device=DEV - ) - output = out_buf[:total].view(batch, tokens, NUM_V_HEADS, HEAD_V_DIM) - elif spec.get("output") == "fp32": - output = torch.zeros(batch, tokens, NUM_V_HEADS, HEAD_V_DIM, device=DEV) - kwargs = dict( q=q, k=k, @@ -164,16 +132,14 @@ def randn(*shape, dtype=torch.float32): initial_state_indices=indices, scale=None, use_qk_l2norm_in_kernel=True, - cu_seqlens=cu_seqlens, + cu_seqlens=None, disable_state_update=mode == "verify", disable_output_calculation=False, intermediate_states_buffer=buffer, cache_steps=cache_steps, - output=output, - ) - return kwargs, dict( - layered=layered, out_buf=out_buf, total=total, live=[s for s in slots if s >= 0] + output=None, ) + return kwargs, list(slots) def _seed(why): @@ -183,55 +149,23 @@ def _seed(why): @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name certifying(op) - kw, ctx = _build(cell.spec, seed=_seed(cell.why)) - pool, buffer, live = kw["initial_state_source"], kw["intermediate_states_buffer"], ctx["live"] - pool_before = None if pool is None else pool.clone() - buffer_before = None if buffer is None else buffer.clone() - layered_before = None if ctx["layered"] is None else ctx["layered"].clone() - tail_before = None if ctx["out_buf"] is None else ctx["out_buf"][ctx["total"] :].clone() + kw, live = _build(cell.spec, seed=_seed(cell.why)) + pool, buffer = kw["initial_state_source"], kw["intermediate_states_buffer"] + pool_before, buffer_before = pool.clone(), buffer.clone() ref = op.reference(**kw) got = op(**kw) _compare((got, pool, buffer), ref) # The in-place locks, bitwise, independent of the reference. - if kw["output"] is not None: - assert got.data_ptr() == kw["output"].data_ptr(), "output must be written in place" - if tail_before is not None: - assert torch.equal(ctx["out_buf"][ctx["total"] :], tail_before), "buffer tail touched" - if pool is not None: - if kw["disable_state_update"]: - assert torch.equal(pool, pool_before), "the pool must not move on a verify call" - else: - for s in live: - assert not torch.equal(pool[s], pool_before[s]), f"slot {s} was not written" - pool_before[live] = pool[live] - assert torch.equal(pool, pool_before), "a slot outside the indices was written" - if buffer is not None: - for s in live: - assert not (buffer[s] == CACHE_SENTINEL).any(), f"row {s} not fully cached" - buffer_before[live] = buffer[live] - assert torch.equal(buffer, buffer_before), "a buffer row outside the indices was written" - if layered_before is not None: - layered = ctx["layered"] - assert torch.equal(layered[0], layered_before[0]) and torch.equal( - layered[2], layered_before[2] - ), "the neighbouring layers must be untouched" - if cell.spec["mode"] == "both": - # Independent of the reference: the last cached step is the stored state. - assert torch.equal(buffer[live, cell.spec["p"] - 1], pool[live]) - - -# ── the mechanism the neighbouring entry fails on ───────────────────────── - - -def test_the_wrapped_kernel_carries_no_autotuner() -> None: - """Why the cold call is safe here: `triton.heuristics` + `triton.jit`, so - there is no trial launch to consume the in-place state.""" - kernel = _impl.fused_recurrent_gated_delta_rule_update_fwd_kernel - names = [cls.__name__ for cls in type(kernel).__mro__] - assert "Autotuner" not in names, f"the wrapped kernel gained an autotuner ({names})" - assert not hasattr(kernel, "configs"), "an autotuner config list appeared on the wrapped kernel" + assert torch.equal(pool, pool_before), "the pool must not move on a verify call" + for s in live: + assert not (buffer[s] == CACHE_SENTINEL).any(), f"row {s} not fully cached" + buffer_before[live] = buffer[live] + assert torch.equal(buffer, buffer_before), "a buffer row outside the indices was written" + + +# ── the mechanism the neighbouring entry failed on ──────────────────────── _COLD_SPEC = dict(mode="write_back", n=2, p=3, slots=(5, 2)) @@ -298,171 +232,7 @@ def test_cold_call_in_a_fresh_process_is_correct() -> None: assert result["out_bitwise_equal"] and result["pool_bitwise_equal"], result -def test_the_gate_rejects_a_doubly_applied_state() -> None: - """The damage an autotuner trial launch leaves -- the state applied twice -- - is rejected by `_compare`, so a cold-call error of that shape cannot pass.""" - kw, _ = _build(_COLD_SPEC, seed=8) - pool = kw["initial_state_source"] - ref = op.reference(**kw) - _compare((op(**kw), pool, None), ref) # one application is right - - doubled = op(**kw) - out_err = (doubled.float() - ref[0].float()).abs().max().item() - pool_err = (pool - ref[1]).abs().max().item() - print(f"\ndoubly applied: output {out_err:.3e}, state {pool_err:.3e}") - assert out_err > 1e-2 and pool_err > 1e-2 - with pytest.raises(AssertionError): - _compare((doubled, pool, None), ref) - - -# ── properties beyond a single reference comparison ─────────────────────── - - -def test_varlen_form_matches_the_batched_form_bitwise() -> None: - """`cu_seqlens` at B = 1 is the same arithmetic as the uniform batch.""" - n, p = 3, 4 - kw, _ = _build(dict(mode="write_back", n=n, p=p, slots=(2, 6, 3)), seed=60) - pool_before = kw["initial_state_source"].clone() - batched = op(**kw) - total = n * p - flat = dict( - kw, - q=kw["q"].reshape(1, total, NUM_K_HEADS, HEAD_K_DIM), - k=kw["k"].reshape(1, total, NUM_K_HEADS, HEAD_K_DIM), - v=kw["v"].reshape(1, total, NUM_V_HEADS, HEAD_V_DIM), - g=kw["g"].reshape(1, total, NUM_V_HEADS), - beta=kw["beta"].reshape(1, total, NUM_V_HEADS), - initial_state_source=pool_before.clone(), - cu_seqlens=torch.arange(0, total + 1, p, device=DEV, dtype=torch.int64), - ) - varlen = op(**flat) - assert torch.equal(varlen.reshape(n, p, NUM_V_HEADS, HEAD_V_DIM), batched) - assert torch.equal(flat["initial_state_source"], kw["initial_state_source"]) - - -def test_scale_default_is_bitwise_head_dim_pow_minus_half() -> None: - kw, _ = _build(dict(mode="write_back", n=2, p=3, slots=(1, 2)), seed=66) - pool_a, pool_b = kw["initial_state_source"].clone(), kw["initial_state_source"].clone() - out_a = op(**dict(kw, initial_state_source=pool_a)) - out_b = op(**dict(kw, initial_state_source=pool_b, scale=HEAD_K_DIM**-0.5)) - assert torch.equal(out_a, out_b) and torch.equal(pool_a, pool_b) - - -def test_deterministic() -> None: - kw, _ = _build(dict(mode="both", n=3, p=4), seed=68) - pool0, buf0 = kw["initial_state_source"].clone(), kw["intermediate_states_buffer"].clone() - runs = [] - for _ in range(2): - pool, buf = pool0.clone(), buf0.clone() - out = op(**dict(kw, initial_state_source=pool, intermediate_states_buffer=buf)) - runs.append((out, pool, buf)) - assert all(torch.equal(x, y) for x, y in zip(*runs)) - - -def test_duplicate_slot_indices_race_on_the_write_back() -> None: - """Two sequences on one slot: the outputs stay right, the slot does not. - Nothing raises, so it is a caller obligation recorded as a measurement.""" - kw, _ = _build(dict(mode="write_back", n=2, p=3, slots=(3, 3)), seed=69) - pool = kw["initial_state_source"] - pool_before = pool.clone() - ref = op.reference(**kw) - got = op(**kw) - torch.testing.assert_close(got, ref[0]) - - def one(i, source): - rows = slice(i, i + 1) - return op.reference( - **dict( - kw, - q=kw["q"][rows], - k=kw["k"][rows], - v=kw["v"][rows], - g=kw["g"][rows], - beta=kw["beta"][rows], - initial_state_source=source, - initial_state_indices=kw["initial_state_indices"][rows], - ) - )[1] - - after_both = one(1, one(0, pool_before)) - drift = (pool[3] - after_both[3]).abs().max().item() - print(f"\nshared slot drift from a sequential application: {drift:.3e}") - assert drift > 1e-2 - - -# ── what the guard refuses, and what the op refuses itself ──────────────── - - -def _small(seed, **spec): - return _build(dict(dict(mode="write_back", n=2, p=3), **spec), seed=seed)[0] - - -def test_a_strided_output_clobbers_its_neighbour_and_is_refused() -> None: - """Unguarded, a last-dim slice of a doubly wide buffer is written with - packed arithmetic: the half never offered is overwritten.""" - kw = _small(70) - wide = torch.zeros(2, 3, NUM_V_HEADS, 2 * HEAD_V_DIM, dtype=torch.bfloat16, device=DEV) - strided = wide[..., :HEAD_V_DIM] - op(**dict(kw, initial_state_source=kw["initial_state_source"].clone(), output=strided)) - assert (wide[..., HEAD_V_DIM:] != 0).any(), "the hazard this guard covers is gone" - with pytest.raises(AssertionError, match="output must be contiguous"): - with validating(op): - op(**dict(kw, output=wide[..., :HEAD_V_DIM])) - - -def test_output_with_the_calculation_off_is_never_written_and_is_refused() -> None: - kw = _build(dict(mode="verify", n=2, p=3), seed=71)[0] - buf = torch.full((2, 3, NUM_V_HEADS, HEAD_V_DIM), -7.0, dtype=torch.bfloat16, device=DEV) - got = op(**dict(kw, disable_output_calculation=True, output=buf)) - assert tuple(got.shape) == (1, 1, 1, 1) and (buf == -7.0).all(), "the hazard is gone" - with pytest.raises(AssertionError, match="mutually exclusive"): - with validating(op): - op(**dict(kw, disable_output_calculation=True, output=buf)) - - -def test_a_strided_pool_loses_the_write_back_and_is_refused() -> None: - """Unguarded, a `[::2]` slot view is replaced by a contiguous copy: the - output is right, and the caller's pool comes back bitwise unchanged.""" - kw = _small(72) - backing = torch.zeros(2 * SLOTS, NUM_V_HEADS, HEAD_V_DIM, HEAD_K_DIM, device=DEV) - strided = backing[::2] - strided.copy_(kw["initial_state_source"]) - before = strided.clone() - ref = op.reference(**dict(kw, initial_state_source=before)) - got = op(**dict(kw, initial_state_source=strided)) - torch.testing.assert_close(got, ref[0]) - assert torch.equal(strided, before), "the hazard this guard covers is gone" - with pytest.raises(AssertionError, match="initial_state_source must be contiguous"): - with validating(op): - op(**dict(kw, initial_state_source=strided)) - - -def test_a_strided_step_buffer_loses_every_cached_state_and_is_refused() -> None: - """The same copy-in trap on the per-step buffer: unguarded, the caller's - buffer comes back bitwise unchanged and nothing raises.""" - kw = _build(dict(mode="verify", n=2, p=3), seed=73)[0] - backing = torch.full( - (2 * SLOTS, 3, NUM_V_HEADS, HEAD_V_DIM, HEAD_K_DIM), CACHE_SENTINEL, device=DEV - ) - strided = backing[::2] - op(**dict(kw, intermediate_states_buffer=strided)) - assert (backing == CACHE_SENTINEL).all(), "the hazard this guard covers is gone" - with pytest.raises(AssertionError, match="intermediate_states_buffer must be contiguous"): - with validating(op): - op(**dict(kw, intermediate_states_buffer=strided)) - - -def test_a_short_cache_steps_writes_into_the_next_row_and_is_refused() -> None: - """Unguarded, cache_steps 2 at P 3: row 0's last step lands in row 1, a - row the call never named.""" - kw = _build(dict(mode="verify", n=1, p=3), seed=74)[0] - short = torch.full((SLOTS, 2, NUM_V_HEADS, HEAD_V_DIM, HEAD_K_DIM), CACHE_SENTINEL, device=DEV) - op(**dict(kw, intermediate_states_buffer=short, cache_steps=2)) - assert not (short[1, 0] == CACHE_SENTINEL).all(), "the hazard this guard covers is gone" - for cache_steps in (2, None): - with pytest.raises(AssertionError, match="cache_steps must be at least"): - with validating(op): - op(**dict(kw, intermediate_states_buffer=short.clone(), cache_steps=cache_steps)) +# ── what the guard refuses ───────────────────────────────────────────── def test_cache_steps_off_the_buffer_step_axis_misplaces_rows_and_is_refused() -> None: @@ -478,103 +248,3 @@ def test_cache_steps_off_the_buffer_step_axis_misplaces_rows_and_is_refused() -> with pytest.raises(AssertionError, match="must equal intermediate_states_buffer's step axis"): with validating(op): op(**dict(kw, intermediate_states_buffer=deep.clone(), cache_steps=2)) - - -def test_op_rejects_what_it_really_rejects() -> None: - """Three loud refusals; the guard admits each so the op's own error is what - surfaces.""" - from triton.compiler.errors import CompilationError - - kw = _small(76) - cu = torch.tensor([0, 3, 6], device=DEV, dtype=torch.int64) - with pytest.raises(ValueError, match="batch size is expected to be 1"): - with validating(op): - op(**dict(kw, cu_seqlens=cu)) - - flat = _build(dict(mode="write_back", n=1, p=6), seed=77)[0] - with pytest.raises(ValueError, match="number of initial states"): - with validating(op): - op(**dict(flat, cu_seqlens=cu)) - - with pytest.raises(CompilationError): - with validating(op): - op(**dict(kw, initial_state_source=None)) - - -def test_sequence_count_past_the_grid_cap_fails_loudly() -> None: - """gridDim.z is `N * HV`, capped at 65535: 1365 sequences at HV 48. Above - it the launch is refused, not truncated, and the context survives.""" - over = 1400 - assert over * NUM_V_HEADS > 65535 - kw = _build(dict(mode="write_back", n=over, p=1, slots=(0,)), seed=78)[0] - indices = torch.arange(over, device=DEV, dtype=torch.int32) % SLOTS - with pytest.raises(RuntimeError, match="invalid argument"): - op(**dict(kw, initial_state_indices=indices, disable_state_update=True)) - - kw, _ = _build(dict(mode="verify", n=2, p=2), seed=79) - ref = op.reference(**kw) - _compare((op(**kw), kw["initial_state_source"], kw["intermediate_states_buffer"]), ref) - - -# ── discrimination ──────────────────────────────────────────────────────── - - -def _wrong_variants(kw, pool_before): - """Plausible wrong recurrences, each expressed as `reference` driven on - transformed inputs -- so the rival is the same arithmetic, misapplied.""" - base = dict(kw, initial_state_source=pool_before) - n, p = kw["q"].shape[:2] - - def flip(t): - return t.flip(1) - - def per_token(t): - return t.reshape((n * p, 1) + tuple(t.shape[2:])) - - out = { - "no_gating": op.reference(**dict(base, g=torch.zeros_like(kw["g"])))[0], - "beta_one": op.reference(**dict(base, beta=torch.ones_like(kw["beta"])))[0], - "qk_swapped": op.reference(**dict(base, q=kw["k"], k=kw["q"]))[0], - "no_initial_state": op.reference( - **dict(base, initial_state_source=torch.zeros_like(pool_before)) - )[0], - "transposed_pool": op.reference( - **dict(base, initial_state_source=pool_before.transpose(-1, -2).contiguous()) - )[0], - # Tiled instead of interleaved: v-head hv reads q/k head hv % H. - "repeat_not_interleave": op.reference( - **dict(base, q=kw["q"].repeat(1, 1, REPEAT, 1), k=kw["k"].repeat(1, 1, REPEAT, 1)) - )[0], - "reversed_tokens": flip( - op.reference(**dict(base, **{x: flip(kw[x]) for x in ("q", "k", "v", "g", "beta")}))[0] - ), - # Every predicted token re-read from the same initial state. - "independent_tokens": op.reference( - **dict( - base, - **{x: per_token(kw[x]) for x in ("q", "k", "v", "g", "beta")}, - initial_state_indices=kw["initial_state_indices"].repeat_interleave(p), - disable_state_update=True, - ) - )[0].reshape(kw["v"].shape), - } - return out - - -def test_default_gate_separates_the_correct_reference_from_wrong_ones() -> None: - """`_compare` runs at torch's default band, so it is only as good as its - ability to reject a near miss. Two of these (`reversed_tokens`, - `independent_tokens`) are wrong only when P > 1 -- this entry's axis.""" - kw, _ = _build(dict(mode="write_back", n=3, p=4, slots=(2, 5, 0)), seed=80) - pool = kw["initial_state_source"] - pool_before = pool.clone() - ref = op.reference(**kw) - got = op(**kw) - _compare((got, pool, None), ref) - own = (got.float() - ref[0].float()).abs().max().item() - for name, wrong in _wrong_variants(kw, pool_before).items(): - gap = (wrong.float() - ref[0].float()).abs().max().item() - print(f"{name}: {gap:.3e} from the reference (the kernel: {own:.3e})") - assert gap > 1e-2, f"{name} is not a distinct reference (gap {gap:.3e})" - with pytest.raises(AssertionError): - _compare((got, pool, None), (wrong, ref[1], None)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_sigmoid_gating_delta_rule_update.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_sigmoid_gating_delta_rule_update.py index 7d0096e9bd80..b1039ed71fd6 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_sigmoid_gating_delta_rule_update.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_sigmoid_gating_delta_rule_update.py @@ -2,11 +2,10 @@ # SPDX-License-Identifier: Apache-2.0 """GPU test for the fused_sigmoid_gating_delta_rule_update entry. -Everything here certifies the vendored Triton branch. The target's decode call -- the +Everything here certifies the vendored Triton branch: the target's decode call -- the rows form (no `cu_seqlens`) over a bf16 pool -- and the packed call over an fp32 pool that -the GDN block test uses as its reference keep the FlashInfer branch out, which -`test_certified_branch_is_the_vendored_triton_kernel` proves. -`test_rows_form_is_batch_invariant` shows the property the bf16 call takes the rows form for. +the GDN block test uses as its reference. `test_rows_form_is_batch_invariant` shows the +property the bf16 call takes the rows form for. """ import zlib @@ -15,7 +14,6 @@ import pytest import torch -import tensorrt_llm._torch.modules.fla.fused_sigmoid_gating_recurrent as _impl from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_sigmoid_gating_delta_rule_update import ( # noqa: E501 ENGINE_SLOTS, ) @@ -65,8 +63,6 @@ def _compare( IN_PROJ_A = MIXED_QKV_WIDTH + NUM_V_HEADS * HEAD_V_DIM # 16384 IN_PROJ_WIDTH = IN_PROJ_A + 2 * NUM_V_HEADS SOFTPLUS_BETA, SOFTPLUS_THRESHOLD = 1.0, 20.0 -_LAYERS = 3 -_TAIL = 5 def _rows_form(mixed, a, b): @@ -80,14 +76,13 @@ def _rows_form(mixed, a, b): def _build(spec, seed): - """`__call__`'s 15 positional arguments, plus the backing buffers some - layouts view. q/k/v are views over one packed `[T, 10240]` conv output, - exactly as the decode call reshapes its column splits.""" + """`__call__`'s 15 positional arguments, plus the packed block the rows form + views. q/k/v are views over one packed `[T, 10240]` conv output, exactly as + the decode call reshapes its column splits.""" g = torch.Generator(device=DEV).manual_seed(seed) seq_lens = spec["seq_lens"] total = sum(seq_lens) - layout = spec.get("layout", "varlen") - if spec.get("ab") == "in_proj": + if spec["ab"] == "in_proj": proj = torch.randn(total, IN_PROJ_WIDTH, generator=g, device=DEV, dtype=torch.bfloat16) mixed = proj[:, :MIXED_QKV_WIDTH] else: @@ -96,54 +91,27 @@ def _build(spec, seed): 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") == "in_proj": + if spec["ab"] == "in_proj": a = proj[:, IN_PROJ_A : IN_PROJ_A + NUM_V_HEADS] b = proj[:, IN_PROJ_A + NUM_V_HEADS :] assert mixed.stride(0) == a.stride(0) == b.stride(0) == IN_PROJ_WIDTH - elif spec.get("ab", "packed") == "separate": + else: a = torch.randn(total, NUM_V_HEADS, generator=g, device=DEV, dtype=torch.bfloat16) b = torch.randn(total, NUM_V_HEADS, generator=g, device=DEV, dtype=torch.bfloat16) - else: - ba = torch.randn(total, 2 * NUM_V_HEADS, generator=g, device=DEV, dtype=torch.bfloat16) - b, a = ba[:, :NUM_V_HEADS], ba[:, NUM_V_HEADS:] A_log = torch.randn(NUM_V_HEADS, generator=g, device=DEV) dt_bias = torch.randn(NUM_V_HEADS, generator=g, device=DEV) - extras = {"mixed": mixed} - slots = spec["slots"] - pool = None - if slots is not None: - lead = (_LAYERS, slots) if layout == "layer_view" else (slots,) - state_shape = (*lead, NUM_V_HEADS, HEAD_V_DIM, HEAD_K_DIM) - pool = torch.randn(state_shape, generator=g, device=DEV) * 0.05 - pool = pool.to(spec.get("pool_dtype", torch.float32)) - if layout == "layer_view": - extras["cache"] = pool - pool = pool[1] + state_shape = (spec["slots"], NUM_V_HEADS, HEAD_V_DIM, HEAD_K_DIM) + pool = torch.randn(state_shape, generator=g, device=DEV) * 0.05 + pool = pool.to(spec.get("pool_dtype", torch.float32)) indices = spec.get("indices") or tuple(range(len(seq_lens))) indices = torch.tensor(indices, dtype=torch.int32, device=DEV) cu_seqlens = torch.tensor( [0] + torch.tensor(seq_lens).cumsum(0).tolist(), dtype=torch.int64, device=DEV ) - if layout == "batched": - n, t = len(seq_lens), seq_lens[0] - q, k = (x.reshape(n, t, NUM_K_HEADS, HEAD_K_DIM).contiguous() for x in (q, k)) - v = v.reshape(n, t, NUM_V_HEADS, HEAD_V_DIM).contiguous() - a, b = (x.reshape(n, t, NUM_V_HEADS).contiguous() for x in (a, b)) - cu_seqlens = None - elif layout == "rows": + if spec.get("layout", "varlen") == "rows": q, k, v, a, b = _rows_form(mixed, a, b) cu_seqlens = None - output = None - if spec.get("output") == "prefix": - extras["buf"] = torch.full( - (1, total + _TAIL, NUM_V_HEADS, HEAD_V_DIM), -7.0, device=DEV, dtype=torch.bfloat16 - ) - output = extras["buf"][:, :total] args = [ A_log, a, @@ -157,69 +125,32 @@ def _build(spec, seed): pool, indices, None, - spec.get("l2norm", True), + True, cu_seqlens, - output, + None, ] - return args, extras + return args, {"mixed": mixed} @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name certifying(op) spec = cell.spec - args, extras = _build(spec, seed=zlib.crc32(cell.why.encode())) + args, _ = _build(spec, seed=zlib.crc32(cell.why.encode())) pool = args[9] - pool_before = None if pool is None else pool.clone() - cache_before = extras["cache"].clone() if "cache" in extras else None - tail_before = extras["buf"][:, -_TAIL:].clone() if "buf" in extras else None + pool_before = pool.clone() ref = op.reference(*args) out = op(*args) _compare((out, pool), ref) - if args[14] is not None: - assert out.data_ptr() == args[14].data_ptr(), "output= must be written in place" - assert torch.equal(extras["buf"][:, -_TAIL:], tail_before), "buffer tail written" - if cache_before is not None: - for other in (0, 2): - assert torch.equal(extras["cache"][other], cache_before[other]), other - if pool is not None: - # The whole-pool compare is the lock on unnamed slots; make it bitwise. - named = {s for s in args[10].tolist() if s >= 0} - unnamed = [s for s in range(pool.shape[0]) if s not in named] - assert torch.equal(pool[unnamed], pool_before[unnamed]), "unnamed slots written" - if spec["slots"] == ENGINE_SLOTS: - for slot in named: - assert not torch.equal(pool[slot], pool_before[slot]), slot - - -def test_certified_branch_is_the_vendored_triton_kernel() -> None: - """The target's rows form and the reference's fp32 packed call keep the FlashInfer arm - out, armed both ways: its entry point is replaced by one that raises; the packed call - over an fp32 pool and the rows form over a bf16 pool complete, and the packed call over a bf16 pool in the same - process does raise -- so the replaced entry point is known to be reachable.""" - - 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) - - rows = dict(seq_lens=(1,) * 32, slots=64, ab="separate", pool_dtype=torch.bfloat16) - args, _ = _build(dict(rows, layout="rows"), seed=102) - assert op(*args).shape == (32, 1, 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 + # The whole-pool compare is the lock on unnamed slots; make it bitwise. + named = {s for s in args[10].tolist() if s >= 0} + unnamed = [s for s in range(pool.shape[0]) if s not in named] + assert torch.equal(pool[unnamed], pool_before[unnamed]), "unnamed slots written" + if spec["slots"] == ENGINE_SLOTS: + for slot in named: + assert not torch.equal(pool[slot], pool_before[slot]), slot #: Batch sizes either side of every FlashInfer dispatch boundary at 48 value heads (its @@ -271,97 +202,6 @@ def test_rows_form_is_batch_invariant() -> None: ) -def test_rows_form_is_the_packed_triton_call_bitwise() -> None: - """At either pool dtype the rows form computes what the packed Triton call computes, - bit for bit -- output and pool. The packed call is forced onto the Triton kernel for - the bf16 comparison.""" - for pool_dtype in (torch.float32, torch.bfloat16): - for n in (1, 8, 32): - spec = dict(seq_lens=(1,) * n, slots=64, ab="separate", pool_dtype=pool_dtype) - packed, _ = _build(spec, seed=300 + n) - rows, _ = _build(dict(spec, layout="rows"), seed=300 + n) - with pytest.MonkeyPatch.context() as mp: - mp.setattr(_impl, "_can_use_flashinfer_gdn_decode", lambda *a, **k: False) - o_packed = op(*packed) - o_rows = op(*rows) - assert torch.equal(o_rows.reshape(o_packed.shape), o_packed), (pool_dtype, n) - assert torch.equal(rows[9], packed[9]), (pool_dtype, n) - - -@pytest.mark.parametrize("n", (1, 32)) -def test_rows_form_cuda_graph_replay_is_bitwise_eager(n) -> None: - """The bf16 rows form captured at batch `n` and replayed on new activations, slots and - pool content, against the same call run eagerly. The kernel takes the batch as a - constexpr, so each batch size compiles on its eager call, before the capture.""" - spec = dict( - seq_lens=(1,) * n, - slots=ENGINE_SLOTS, - ab="separate", - layout="rows", - pool_dtype=torch.bfloat16, - ) - args, extras = _build(spec, seed=400 + n) - static = list(args) - for _ in range(2): - op(*static) - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - out_static = op(*static) - for trial in range(3): - slots = tuple(7 * trial + 13 * i for i in range(n)) - fresh, fresh_extras = _build(dict(spec, indices=slots), seed=500 + 10 * n + trial) - for position in (0, 1, 2, 8, 9, 10): - static[position].copy_(fresh[position]) - extras["mixed"].copy_(fresh_extras["mixed"]) - pool_before = static[9].clone() - graph.replay() - torch.cuda.synchronize() - replayed, replayed_pool = out_static.clone(), static[9].clone() - eager_args = list(fresh) - eager_args[9] = pool_before - eager = op(*eager_args) - assert torch.equal(replayed, eager), (n, trial) - assert torch.equal(replayed_pool, pool_before), (n, trial) - - -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_guard_refuses_rows_views_whose_token_stride_is_not_the_row_stride() -> None: """Unguarded, q/k/v reshaped to `[N, 1, heads, 128]` from column slices of the packed block carry `stride(1) = heads * 128`: without `cu_seqlens` the kernel reads row n at @@ -389,38 +229,3 @@ def test_guard_refuses_rows_views_whose_token_stride_is_not_the_row_stride() -> with pytest.raises(AssertionError, match=r"stride\(0\) must equal"): 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_gdn_block.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_block.py index 017b2289cfc0..b551a1ec8097 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_block.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_block.py @@ -4,8 +4,8 @@ Drives `DecodeTarget._gdn_block` and `PrefillTarget._gdn_block` at the checkpoint's geometry (hidden 2560, 16 k-heads, 48 v-heads, head dim 128, -fp32 recurrent pool of 461 slots) on decode, prefill-only and mixed -(context + one-token generation rows) batches, and compares the block output +fp32 recurrent pool of 461 slots) on a decode and a mixed (context + one-token +generation rows) batch, and compares the block output and both state pools with a reference that composes the Triton recurrences (`fused_sigmoid_gating_delta_rule_update` for decode, `chunk_gated_delta_rule` over every row of a prefill step) around the same @@ -354,20 +354,10 @@ def test_decode_block() -> None: ) -#: (context lengths, generation rows, has_initial_state per context): an 8K prefill-only -#: step, and a mixed step of three variable-length contexts (two carrying state) plus -#: generation rows. -PREFILL_CASES = ( - ((8192,), 0, (False,)), - ((200, 37, 91), 5, (True, False, True)), -) - - -@pytest.mark.parametrize( - "case", PREFILL_CASES, ids=lambda c: f"{'+'.join(map(str, c[0]))}+{c[1]}x1" -) -def test_prefill_block(case, warm_triton_chunk) -> None: # noqa: F811 — fixture by name - ctx_lens, n_dec, has_init = case +def test_prefill_block(warm_triton_chunk) -> None: # noqa: F811 — fixture by name + """A mixed step: three variable-length contexts (two carrying state) plus five + one-token generation rows.""" + ctx_lens, n_dec, has_init = (200, 37, 91), 5, (True, False, True) tokens = sum(ctx_lens) + n_dec core, L, x = _core(400 + tokens), _layer(500 + tokens), _x(tokens, 600 + tokens) core_ref = _clone_core(core) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_gated_norm_pdl.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_gated_norm_pdl.py index 3b95ef012e24..24ebe8c685a8 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_gated_norm_pdl.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_gated_norm_pdl.py @@ -5,27 +5,18 @@ Cells against the composition's fp64 reference (the gated delta rule's recurrence, then the gated norm entry's) in torch's default band for the bf16 output and the fp32 pool, the pool compared whole so slots the call did not name are -gated too. Beyond the cells: agreement with the unfused composition it replaces -(FlashInfer's recurrence, then the gated norm entry), PDL on and off bitwise, a -stale-read check under graph replay (a producer rewrites the projection inside the -graph right before the op), and the guards. +gated too. Beyond the cells: a stale-read check under graph replay (a producer +rewrites the projection inside the graph right before the op), and the row guard. """ import zlib -import flashinfer.gdn_decode 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, -) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_decode_gated_norm_pdl import ( # noqa: E501 gdn_decode_gated_norm_pdl as op, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_decode_gated_norm_pdl import ( # noqa: E501 - launch, -) __extra_import_path__ = ["../../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name @@ -59,9 +50,6 @@ def _build(spec, seed): pool = page.as_strided((slots, HV, V, K), (layers * HV * V * K, V * K, K, 1)) for s in (100, 101): pool[s].copy_(torch.randn(HV, V, K, generator=g, device=DEV) * 0.05) - elif spec.get("pool") == "layer_view": - cache = torch.randn(3, SLOTS, HV, V, K, generator=g, device=DEV) * 0.05 - pool = cache[1] else: pool = torch.randn(SLOTS, HV, V, K, generator=g, device=DEV) * 0.05 if spec.get("pool") == "coalesced": @@ -94,54 +82,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by assert (out[1] == 0).all(), "a pad row's output is zero" -def _unfused(proj, pool, slots, A_log, dt_bias, w, eps, h, k): - rows = proj.shape[0] - q, kk, v, z, a, b = torch.split(proj, [h * k, h * k, HV * V, HV * V, HV, HV], 1) - o = torch.empty(rows, 1, HV, V, dtype=torch.bfloat16, device=DEV) - flashinfer.gdn_decode.gated_delta_rule_decode_pretranspose( - q=q.reshape(rows, 1, h, k), - k=kk.reshape(rows, 1, h, k), - v=v.reshape(rows, 1, HV, V), - state=None, - A_log=A_log, - a=a.reshape(rows, 1, HV), - dt_bias=dt_bias, - b=b.reshape(rows, 1, HV), - scale=None, - output=o, - use_qk_l2norm=True, - initial_state=pool, - initial_state_indices=slots, - output_state_indices=None, - ) - return rms_norm_gated_token_major( - o.reshape(-1, V), z.reshape(rows, HV, V), w, eps, None, "sigmoid" - ).reshape(rows, HV * V) - - -def test_matches_the_unfused_composition() -> None: - rows = 4 # MAX_ROWS - args = _build(dict(rows=rows), seed=900 + rows) - pool_unfused = args[1].clone() - want = _unfused(args[0], pool_unfused, *args[2:]) - got = op(*args) - torch.testing.assert_close(got, want) - torch.testing.assert_close(args[1], pool_unfused) - print( - f"\n rows {rows}: output bitwise the composition's {(got == want).float().mean().item():.4f}" - ) - - -def test_pdl_on_and_off_agree_bitwise() -> None: - args = _build(dict(rows=2), seed=31) - pool0 = args[1].clone() - a = launch(*args, pdl=True) - pool_a = args[1].clone() - args[1].copy_(pool0) - b = launch(*args, pdl=False) - assert torch.equal(a, b) and torch.equal(pool_a, args[1]) - - def test_reads_after_the_wait_under_graph_replay() -> None: """A producer rewrites the projection inside the graph right before the op; with PDL a read of it before `griddepcontrol.wait` would see the previous replay's row. The @@ -176,8 +116,7 @@ def step(): def test_guards() -> None: + """More rows than `MAX_ROWS` are refused.""" args = _build(dict(rows=1), seed=5) with validating(op), pytest.raises(AssertionError): op(args[0].expand(5, -1).contiguous(), args[1], args[2].expand(5).contiguous(), *args[3:]) - with validating(op), pytest.raises(AssertionError): - op(args[0], args[1].to(torch.bfloat16), *args[2:]) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py index 77756f9242aa..adf92c5b6ba6 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py @@ -5,30 +5,22 @@ Cells against the composition's fp64 reference (the gated delta rule's recurrence, then the gated norm entry's) in torch's default band for the bf16 output and the fp32 pool, the pool compared whole so slots the call did not name are -gated too. Beyond the cells: agreement with the unfused composition it replaces -(FlashInfer's recurrence, then the gated norm entry) at the 128-row decode batch, PDL -on and off bitwise, a stale-read check under graph replay (a producer rewrites the -projection inside the graph right before the op), and the guards. +gated too. Beyond the cells: a stale-read check under graph replay (a producer +rewrites the projection inside the graph right before the op), and the pool-dtype +guard. The `gdn_decode_rows` entry (the recurrence alone) has its own cells below. """ import zlib -import flashinfer.gdn_decode 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, -) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_decode_rows_gated_norm import ( # noqa: E501 gdn_decode_rows as rows_op, ) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_decode_rows_gated_norm import ( # noqa: E501 gdn_decode_rows_gated_norm as op, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_decode_rows_gated_norm import ( # noqa: E501 - launch, -) __extra_import_path__ = ["../../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name @@ -62,9 +54,6 @@ def _build(spec, seed): pool = page.as_strided((slots, HV, V, K), (layers * HV * V * K, V * K, K, 1)) for s in (100, 101): pool[s].copy_(torch.randn(HV, V, K, generator=g, device=DEV) * 0.05) - elif spec.get("pool") == "layer_view": - cache = torch.randn(3, SLOTS, HV, V, K, generator=g, device=DEV) * 0.05 - pool = cache[1] else: pool = torch.randn(SLOTS, HV, V, K, generator=g, device=DEV) * 0.05 if spec.get("pool") == "coalesced": @@ -97,54 +86,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by assert (out[1] == 0).all(), "a pad row's output is zero" -def _unfused(proj, pool, slots, A_log, dt_bias, w, eps, h, k): - rows = proj.shape[0] - q, kk, v, z, a, b = torch.split(proj, [h * k, h * k, HV * V, HV * V, HV, HV], 1) - o = torch.empty(rows, 1, HV, V, dtype=torch.bfloat16, device=DEV) - flashinfer.gdn_decode.gated_delta_rule_decode_pretranspose( - q=q.reshape(rows, 1, h, k), - k=kk.reshape(rows, 1, h, k), - v=v.reshape(rows, 1, HV, V), - state=None, - A_log=A_log, - a=a.reshape(rows, 1, HV), - dt_bias=dt_bias, - b=b.reshape(rows, 1, HV), - scale=None, - output=o, - use_qk_l2norm=True, - initial_state=pool, - initial_state_indices=slots, - output_state_indices=None, - ) - return rms_norm_gated_token_major( - o.reshape(-1, V), z.reshape(rows, HV, V), w, eps, None, "sigmoid" - ).reshape(rows, HV * V) - - -def test_matches_the_unfused_composition() -> None: - rows = 128 - args = _build(dict(rows=rows), seed=900 + rows) - pool_unfused = args[1].clone() - want = _unfused(args[0], pool_unfused, *args[2:]) - got = op(*args) - torch.testing.assert_close(got, want) - torch.testing.assert_close(args[1], pool_unfused) - print( - f"\n rows {rows}: output bitwise the composition's {(got == want).float().mean().item():.4f}" - ) - - -def test_pdl_on_and_off_agree_bitwise() -> None: - args = _build(dict(rows=8), seed=31) - pool0 = args[1].clone() - a = launch(*args, pdl=True) - pool_a = args[1].clone() - args[1].copy_(pool0) - b = launch(*args, pdl=False) - assert torch.equal(a, b) and torch.equal(pool_a, args[1]) - - def test_reads_after_the_wait_under_graph_replay() -> None: """A producer rewrites the projection inside the graph right before the op; with PDL a read of it before `griddepcontrol.wait` would see the previous replay's row. The @@ -177,13 +118,10 @@ def step(): def test_guards() -> None: + """A bf16 pool (the target sends this entry an fp32 one only) is refused.""" args = _build(dict(rows=4), seed=5) with validating(op), pytest.raises(AssertionError): op(args[0], args[1].to(torch.bfloat16), *args[2:]) - with validating(op), pytest.raises(AssertionError): - op(args[0], args[1], args[2].long(), *args[3:]) - with validating(op), pytest.raises(AssertionError): - op(args[0][:, :-2], *args[1:]) # ---- gdn_decode_rows: the recurrence alone, into a row range of a step's output ---------- @@ -215,21 +153,3 @@ def test_rows_certified_cells(cell, certifying) -> None: # noqa: F811 — fixtu ) if cell.spec.get("pad"): assert (got[1] == 0).all(), "a pad row's output is zero" - - -def test_rows_then_gated_norm_matches_the_fused_entry() -> None: - """The recurrence alone, then the gated norm entry: the fused entry's output within its - band, and the same pool bitwise (one kernel's recurrence), at a mixed step's 121 - generation rows.""" - rows = 121 - args = _build(dict(rows=rows), seed=700 + rows) - pool_fused = args[1].clone() - fused = op(args[0], pool_fused, *args[2:]) - o = torch.empty(rows, HV * V, dtype=torch.bfloat16, device=DEV) - rows_op(args[0], args[1], args[2], args[3], args[4], args[7], args[8], o) - assert torch.equal(args[1], pool_fused), "pool" - z = torch.split(args[0], [2 * H * K, HV * V, HV * V, 2 * HV], 1)[2] - normed = rms_norm_gated_token_major( - o.reshape(-1, V), z.reshape(rows, HV, V), args[5], args[6], None, "sigmoid" - ).reshape(rows, HV * V) - torch.testing.assert_close(normed, fused) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_in_proj_conv1d_pdl.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_in_proj_conv1d_pdl.py index dcaf4fb3c133..e6a5f55bd730 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_in_proj_conv1d_pdl.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_in_proj_conv1d_pdl.py @@ -5,10 +5,9 @@ Cells against the entry's fp32 reference: the output within the band the decode projections use (4 bf16 ulp element-wise against the row scale, 2 ulp relative RMS -- the GEMV's summation order is the only difference), the pool -the same, a pad row's slot bitwise untouched. Beyond the cells: agreement with -the unfused composition it replaces (`low_m_bf16_mm` + `causal_conv1d_update`); -two decode steps under CUDA-graph replay, the second reading the window the -first one wrote; PDL on and off bitwise; and the guards. +the same, a pad row's slot bitwise untouched. Beyond the cells: two decode +steps under CUDA-graph replay, the second reading the window the first one +wrote; and the row guard. """ import zlib @@ -16,12 +15,6 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.gemm.low_m_bf16_mm import ( # noqa: E501 - low_m_bf16_mm, -) -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, -) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_in_proj_conv1d_pdl import ( # noqa: E501 MAX_ROWS, ) @@ -73,27 +66,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 assert torch.equal(pool[slot], state[slot]), f"slot {slot} was written" -def test_matches_the_unfused_composition() -> None: - rows = MAX_ROWS - x, mat_b, state, taps = _build(rows, 461, seed=5 + rows) - slots = torch.tensor([37, 211][:rows], dtype=torch.int32, device="cuda") - fused_pool, unfused_pool = state.clone(), state.clone() - fused = op(x, mat_b, fused_pool, taps, slots, CONV_DIM, PAD) - unfused = low_m_bf16_mm(x, mat_b, False) - causal_conv1d_update( - x=unfused[:, :CONV_DIM], - conv_state=unfused_pool, - weight=taps, - bias=None, - silu_activation=True, - cache_seqlens=None, - conv_state_indices=slots, - pad_slot_id=PAD, - ) - _compare(fused, unfused) - _compare(fused_pool, unfused_pool) - - def test_two_steps_under_graph_replay() -> None: """Decode runs this under CUDA graphs with PDL, layer after layer: capture one step, replay it twice with new activations. The second replay's conv reads the @@ -123,33 +95,9 @@ def test_two_steps_under_graph_replay() -> None: _compare(pool, want_state) -def test_pdl_on_and_off_agree_bitwise() -> None: - x, mat_b, state, taps = _build(2, 461, seed=13) - slots = torch.tensor([3, 400], dtype=torch.int32, device="cuda") - on_pool, off_pool = state.clone(), state.clone() - on = op(x, mat_b, on_pool, taps, slots, CONV_DIM, PAD, True) - off = op(x, mat_b, off_pool, taps, slots, CONV_DIM, PAD, False) - assert torch.equal(on, off) and torch.equal(on_pool, off_pool) - - -def test_band_rejects_a_conv_on_the_wrong_columns() -> None: - """The near miss a column offset would produce: the conv applied one column off.""" - x, mat_b, state, taps = _build(1, 461, seed=17) - slots = torch.tensor([42], dtype=torch.int32, device="cuda") - out = op(x, mat_b, state.clone(), taps, slots, CONV_DIM, PAD) - ref_out, _ = op.reference(x, mat_b, state, taps, slots, CONV_DIM, PAD) - _compare(out, ref_out) - with pytest.raises(AssertionError): - _compare(out, torch.roll(ref_out, 1, dims=1)) - - def test_guards() -> None: + """More rows than `MAX_ROWS` are refused.""" x, mat_b, state, taps = _build(MAX_ROWS + 1, 16, seed=21) slots = torch.zeros(MAX_ROWS + 1, dtype=torch.int32, device="cuda") with validating(op), pytest.raises(AssertionError): op(x, mat_b, state, taps, slots, CONV_DIM, PAD) - x1 = x[:1].contiguous() - with validating(op), pytest.raises(AssertionError): - op(x1, mat_b.contiguous(), state, taps, slots[:1], CONV_DIM, PAD) - with validating(op), pytest.raises(AssertionError): - op(x1, mat_b, state, taps, slots[:1].long(), CONV_DIM, PAD) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_prefill_conv_post.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_prefill_conv_post.py index e9ce6b8488e9..8482175c4a3a 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_prefill_conv_post.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_prefill_conv_post.py @@ -3,11 +3,7 @@ """GPU test for the gdn_prefill_conv_post entry. Every cell is checked against the entry's fp32/fp64 reference, the conv pool -bitwise. The three-call composition it replaces in the target -(`extract_transpose_prefill_slice` -> `causal_conv1d_fwd` -> -`fused_gdn_post_conv`, each certified on its own) is checked once, at the -target-call cell, again with the conv pool compared bitwise. -Geometry is a Qwen3.8-Flash-Next linear-attention layer: 16 x 128 key heads, +bitwise. Geometry is a Qwen3.8-Flash-Next linear-attention layer: 16 x 128 key heads, 48 x 128 value heads, qkv_dim 10240, conv width 4, bf16 activations. """ @@ -17,15 +13,6 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.causal_conv1d_fwd import ( # noqa: E501 - causal_conv1d_fwd, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.extract_transpose_prefill_slice import ( # noqa: E501 - extract_transpose_prefill_slice, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_gdn_post_conv import ( # noqa: E501 - fused_gdn_post_conv, -) from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_prefill_conv_post import ( # noqa: E501 A_LOG_RANGE, CONV_WIDTH, @@ -93,8 +80,6 @@ def randn(*shape): A_log = torch.rand(NUM_V_HEADS, generator=g, device=DEV) * (hi_ - lo) + lo lo, hi_ = DT_BIAS_RANGE dt_bias = torch.rand(NUM_V_HEADS, generator=g, device=DEV) * (hi_ - lo) + lo - if "dt_bias_fill" in spec: - dt_bias = torch.full((NUM_V_HEADS,), spec["dt_bias_fill"], device=DEV) return [ x, weight, @@ -113,19 +98,6 @@ def randn(*shape): ] -def _unfused(args): - """The target's previous three calls on a copy of the pool: `(q, k, v, g, - beta, pool)`.""" - x, weight, conv_states, qsl, si, hi, a, b, A_log, dt_bias = args[:10] - pool = conv_states.clone() - xt = extract_transpose_prefill_slice(x, x.shape[0], 0, QKV_DIM) - causal_conv1d_fwd(xt, weight, None, pool, qsl, si, hi, True, PAD, None) - out = fused_gdn_post_conv( - xt, None, a, b, A_log, dt_bias, *HEADS, L2_NORM_EPS, SOFTPLUS_THRESHOLD - ) - return (*out, pool) - - def _compare(out, pool, want, what=""): q, k, v, g, beta = out q_r, k_r, v_r, g_r, beta_r, pool_r = want @@ -148,7 +120,6 @@ def _compare(out, pool, want, what=""): 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())) - pool0 = args[2].clone() want = op.reference(*args) inputs = [t.clone() if isinstance(t, torch.Tensor) else t for t in args] out = op(*args) @@ -156,38 +127,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res if isinstance(arg, torch.Tensor) and i != 2: assert torch.equal(arg, before), f"input {i} must be read-only" _compare(out, args[2], want, what=f"{cell.why[:30]} vs reference") - if PAD in args[4].tolist(): - pad_free = [s for s in range(pool0.shape[0]) if s not in args[4].tolist()] - assert torch.equal(args[2][pad_free], pool0[pad_free]), "unnamed slots written" - - -def test_matches_the_unfused_composition() -> None: - """The target-call cell against the three calls the entry replaces; each of - those is certified on its own, so one representative call is the check.""" - cell = next(c for c in op.CELLS if c.why.startswith("target call")) - args = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) - old = _unfused(args) - out = op(*args) - _compare(out, args[2], old, what=f"{cell.why[:30]} vs unfused") - - -def test_carry_and_silu_are_load_bearing() -> None: - """The band separates a dropped carry and a dropped SiLU.""" - spec = dict(lengths=(5, 70), pool_slots=4, slots=(1, 3), initial=(True, True)) - args = _build(spec, seed=301) - want = op.reference(*args) - no_carry = list(args) - no_carry[5] = torch.zeros_like(args[5]) - out = op(*[t.clone() if isinstance(t, torch.Tensor) else t for t in no_carry]) - with pytest.raises(AssertionError): - torch.testing.assert_close(out[2], want[2]) - x, weight = args[0], args[1] - # No SiLU: the pre-activation of the second context's last token, v channels. - pre = sum(x[74 - 3 + j, :QKV_DIM].float() * weight[:, j].float() for j in range(4)) - pre = pre[2 * NUM_K_HEADS * HEAD_K_DIM :] - got = op(*_build(spec, seed=301))[2][0, 74].reshape(-1).float() - assert (got - torch.nn.functional.silu(pre)).abs().max() < 0.05 - assert (got - pre).abs().max() > 0.5 def test_guard_refuses_channel_strided_x() -> None: @@ -197,18 +136,3 @@ def test_guard_refuses_channel_strided_x() -> None: with pytest.raises(AssertionError, match="unit channel stride"): with validating(op): op(*bad) - - -def test_guard_refuses_short_a_b_and_strided_gates() -> None: - args = _build(dict(lengths=(16, 8), pool_slots=2, slots=(0, 1), initial=(True, True)), seed=303) - for position in (6, 7): - bad = list(args) - bad[position] = args[position][:16] - with pytest.raises(AssertionError, match="at least T=24 rows"): - with validating(op): - op(*bad) - bad = list(args) - bad[8] = torch.stack([args[8], args[8]], dim=1)[:, 0] - with pytest.raises(AssertionError, match="contiguous 1-D"): - with validating(op): - op(*bad) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_pack_gdn_decode_qkv.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_pack_gdn_decode_qkv.py index 50a569172aac..4adc8c137f6f 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_pack_gdn_decode_qkv.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_pack_gdn_decode_qkv.py @@ -2,15 +2,13 @@ # SPDX-License-Identifier: Apache-2.0 """GPU test for the pack_gdn_decode_qkv entry. -The op is a pure copy, so every comparison is bitwise; the discrimination test -is what shows that gate is not vacuous on these operands. +The op is a pure copy, so every comparison is bitwise. """ import zlib import pytest import torch -from triton.compiler.errors import CompilationError from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.pack_gdn_decode_qkv import ( # noqa: E501 CKPT_HEADS, @@ -99,52 +97,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res 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") - - -def test_op_rejects_head_dims_that_round_to_different_blocks() -> None: - """Loud at compile time, before any launch -- so not an `is_valid` guard.""" - heads = (16, 128, 48, 64) - 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. @@ -163,13 +115,3 @@ def test_a_short_row_reads_the_next_row_and_is_refused() -> None: with pytest.raises(AssertionError, match="must be at least qkv_dim=10240 wide"): with validating(op): op(mixed, *CKPT_HEADS) - - -def test_guard_refuses_a_three_dimensional_input() -> None: - """`[batch, draft, qkv_dim]` -- a caller who skipped the flatten. Unguarded - it faults with an illegal memory access (measured; not replayed here, it - poisons the process); the guard refuses it before the launch.""" - mixed = torch.randn(4, 2, QKV_DIM, device=DEV).to(torch.bfloat16) - with pytest.raises(AssertionError, match="must be at least qkv_dim="): - with validating(op): - op(mixed, *CKPT_HEADS) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_reset_recurrent_state_rows.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_reset_recurrent_state_rows.py index 55c78cddac33..13a81044a508 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_reset_recurrent_state_rows.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_reset_recurrent_state_rows.py @@ -47,38 +47,25 @@ def _compare( def _build(spec, seed): - """`__call__`'s arguments, plus the backing allocations a layout views.""" + """`__call__`'s arguments, plus the backing allocations a layer view views.""" g = torch.Generator(device=DEV).manual_seed(seed) slots = spec["slots"] - rec_shape, conv_shape = REC_SHAPE, CONV_SHAPE rec_dtype = spec.get("rec_dtype", torch.float32) - layout = spec.get("layout", "plain") + layer_view = spec.get("layout", "plain") == "layer_view" def randn(*shape, dtype): return torch.randn(*shape, generator=g, device=DEV).to(dtype) backing = {} - if layout == "row_padding": - row = 1 - for d in rec_shape: - row *= d - backing["rec"] = randn(slots, 2, row, dtype=rec_dtype) - rec = backing["rec"][:, 0] - assert rec.stride() == (2 * row, 1) - elif layout == "layer_view": - backing["rec"] = randn(_LAYERS, slots, *rec_shape, dtype=rec_dtype) - rec = backing["rec"][1] + if layer_view: + backing["rec"] = randn(_LAYERS, slots, *REC_SHAPE, dtype=rec_dtype) + backing["conv"] = randn(_LAYERS, slots, *CONV_SHAPE, dtype=torch.bfloat16) + rec, conv = backing["rec"][1], backing["conv"][1] else: - rec = randn(slots, *rec_shape, dtype=rec_dtype) - conv = None - if spec.get("conv", True): - if layout == "layer_view": - backing["conv"] = randn(_LAYERS, slots, *conv_shape, dtype=torch.bfloat16) - conv = backing["conv"][1] - else: - conv = randn(slots, *conv_shape, dtype=torch.bfloat16) - - indices = torch.tensor(spec["indices"], dtype=spec.get("index_dtype", torch.int32), device=DEV) + rec = randn(slots, *REC_SHAPE, dtype=rec_dtype) + conv = randn(slots, *CONV_SHAPE, dtype=torch.bfloat16) + + indices = torch.tensor(spec["indices"], dtype=torch.int32, device=DEV) flags = torch.tensor(spec["flags"], dtype=torch.bool, device=DEV) return [rec, indices, flags, conv], backing @@ -105,54 +92,12 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res # Non-vacuity: the engine-width cells reset at least one row, so the # whole-pool compare is also the lock on the ~430 unnamed slots. assert reset - if spec.get("layout") == "row_padding": - assert torch.equal(backing["rec"][:, 1], before["rec"][:, 1]), "padding lane written" if spec.get("layout") == "layer_view": for name in backing: for other in (0, 2): assert torch.equal(backing[name][other], before[name][other]), (name, other) -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", - [ - ("slot_count_mismatch", "recurrent and convolution pools must have the same slot count"), - ("recurrent_rows_noncontiguous", "recurrent state rows must be contiguous"), - ("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. One case - per family the callable checks: slot counts, row layout, device.""" - 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 == "slot_count_mismatch": - conv = torch.randn(3, 4, dtype=torch.bfloat16, device=DEV) - elif case == "recurrent_rows_noncontiguous": - rec = torch.randn(2, 8, dtype=torch.float32, device=DEV)[:, ::2] - 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) @@ -160,51 +105,13 @@ def _pools(): return rec, conv -def test_guard_hazard_strided_index_tensors() -> None: - """Unguarded, a stride-2 view is read through its base pointer: indices - [0, 1, 2] reset [0, 1], and flags [T, T, T] reset row 1. No error.""" - good_indices = torch.tensor([0, 1, 2], dtype=torch.int32, device=DEV) - good_flags = torch.zeros(3, dtype=torch.bool, device=DEV) - strided_indices = torch.tensor([0, 9, 1, 9, 2, 9], dtype=torch.int32, device=DEV)[::2] - strided_flags = torch.tensor([True, False] * 3, device=DEV)[::2] - for indices, flags in ((strided_indices, good_flags), (good_indices, strided_flags)): - rec, conv = _pools() - # The reference reads the tensors by value, which is what a caller meant. - want = op.reference(rec, indices, flags, conv) - got = op(rec, indices, flags, conv) - with pytest.raises(AssertionError, match="pool differs"): - _compare(got, want) - - -def test_guard_hazard_two_d_index_tensors() -> None: - """Unguarded, a matching [2, 3] pair passes the op's shape check and only - its first shape[0] = 2 of 6 requests are considered.""" +def test_guard_refuses_index_tensors_the_op_misreads() -> None: + """Unguarded, a stride-2 `state_indices` view is read through its base pointer + as a flat buffer: indices [0, 1, 2] arrive as [0, 9, 1], so rows 0 and 1 are reset + and row 2 is not (measured). The guard refuses it.""" rec, conv = _pools() - indices = torch.arange(6, dtype=torch.int32, device=DEV).reshape(2, 3) - flags = torch.zeros(2, 3, dtype=torch.bool, device=DEV) - want = op.reference(rec, indices.reshape(-1), flags.reshape(-1), conv) - got = op(rec, indices, flags, conv) - with pytest.raises(AssertionError, match="pool differs"): - _compare(got, want) - - -@pytest.mark.parametrize( - "which,how", [("state_indices", "strided"), ("has_initial_states", "two_d")] -) -def test_guard_refuses_index_tensors_the_op_misreads(which: str, how: str) -> None: - rec, conv = _pools() - indices = torch.zeros(3, dtype=torch.int32, device=DEV) + indices = torch.tensor([0, 9, 1, 9, 2, 9], dtype=torch.int32, device=DEV)[::2] flags = torch.zeros(3, dtype=torch.bool, device=DEV) - if how == "strided": - bad_indices = torch.tensor([0, 9, 1, 9, 2, 9], dtype=torch.int32, device=DEV)[::2] - bad_flags = torch.zeros(6, dtype=torch.bool, device=DEV)[::2] - else: - bad_indices = torch.arange(6, dtype=torch.int32, device=DEV).reshape(2, 3) - bad_flags = torch.zeros(2, 3, dtype=torch.bool, device=DEV) - if which == "state_indices": - indices = bad_indices - else: - flags = bad_flags - with pytest.raises(AssertionError, match=f"{which} must be a contiguous 1-D tensor"): + with pytest.raises(AssertionError, match="state_indices must be a contiguous 1-D tensor"): with validating(op): op(rec, indices, flags, conv) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_gdn_replay_verify.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_gdn_replay_verify.py index cef8c8358a46..6e5b42957ea4 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_gdn_replay_verify.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_gdn_replay_verify.py @@ -242,15 +242,13 @@ def _state(gen, cpu, n, t, gating=None): @pytest.mark.parametrize( "n,t,accept", [ - # Below the partition threshold: the history wraps in-kernel. - (8, 4, "random"), - # At the threshold, at max_draft_len 7 (T = 8, the replay kernel's ceiling). + # At the partition threshold, at max_draft_len 7 (T = 8, the replay kernel's + # ceiling): the history wraps through the all-layer commit. (CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE, 8, "random"), - # Every draft accepted: from a fresh slot at T = 8 the second step appends exactly up - # to the history's end (`pnat + T == 16`) and every second step after it wraps -- - # in-kernel below the partition threshold, through the all-layer commit at and above it. + # Below the threshold, every draft accepted: from a fresh slot at T = 8 the second + # step appends exactly up to the history's end (`pnat + T == 16`) and every second + # step after it wraps in-kernel. (4, 8, "full"), - (24, 8, "full"), ], ids=lambda v: str(v), ) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_index_share.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_index_share.py index 357e263bc0c7..548008dacdf4 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_index_share.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_index_share.py @@ -149,7 +149,7 @@ def test_reuse_of_a_canonical_capture_stays_ascending(): """A canonical capture plus the groups completed since is still ascending: appended ids are newer than every captured one, and a full row gives up its oldest ids to them, so the newest captured groups -- the ones next to the new ones -- are kept.""" - for fresh in (0, 1, 2): + for fresh in (0, 2): args, captured_visible = _reuse_case(fresh) out = reuse(**args) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_kv_range_probe.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_kv_range_probe.py index 83cb86883b0e..2603aaf63e00 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_kv_range_probe.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_kv_range_probe.py @@ -42,8 +42,9 @@ def test_a_forward_within_the_e4m3_range_is_silent(monkeypatch): logger.warning.assert_not_called() -@pytest.mark.parametrize("peak", [1000.0, float("nan")]) +@pytest.mark.parametrize("peak", [-600.0, float("nan")]) def test_a_magnitude_past_the_e4m3_maximum_is_reported_once(monkeypatch, peak): + """A negative value counts by its magnitude; a NaN counts as out of range.""" logger = Mock() monkeypatch.setattr(m, "logger", logger) @@ -54,15 +55,6 @@ def test_a_magnitude_past_the_e4m3_maximum_is_reported_once(monkeypatch, peak): assert "e4m3 maximum 448" in logger.warning.call_args.args[0] -def test_a_negative_value_counts_by_magnitude(monkeypatch): - logger = Mock() - monkeypatch.setattr(m, "logger", logger) - - assert _run_forward(m._KvRangeProbe(LAST), (1.0, -600.0, 1.0)) == [True, True, False] - - logger.warning.assert_called_once() - - @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") def test_observing_under_capture_reads_nothing(monkeypatch): logger = Mock() diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_spec_rollback.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_spec_rollback.py index 25df5f9bc522..c56c098ff397 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_spec_rollback.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_spec_rollback.py @@ -258,7 +258,8 @@ def _assert_state(b, expected, what): # (ctx_lens, gen_starts, p, gen_accepted). Positions straddle pages and land on -# compression anchors; one request accepts everything. +# compression anchors; one request accepts everything. `decode` and `mixed` are the accepted- +# prefix check; `decode_p8` drives the shared-page and capture tests. CASES = { "decode": ((), (5, 30, 61, 100), 4, (1, 2, 4, 3)), "mixed": ((7, 3), (12, 29, 95), 4, (2, 1, 3)), @@ -272,7 +273,7 @@ def _slots(num): return list(range(3, 3 + 2 * num, 2)) -@pytest.mark.parametrize("case", sorted(CASES)) +@pytest.mark.parametrize("case", ["decode", "mixed"]) def test_commit_matches_the_accepted_prefix(case): ctx_lens, gen_starts, p, gen_accepted = CASES[case] b = _Batch(11, ctx_lens, gen_starts, p, _slots(len(ctx_lens) + len(gen_starts))) @@ -296,15 +297,6 @@ def test_commit_over_the_strided_planes_of_a_shared_page(): _assert_state(b, _reference(b), case) -@pytest.mark.parametrize("accepted", ["all", "one"]) -def test_commit_at_the_acceptance_extremes(accepted): - ctx_lens, gen_starts, p, _ = CASES["mixed"] - b = _Batch(12, ctx_lens, gen_starts, p, _slots(len(ctx_lens) + len(gen_starts))) - b.set_accepted([p if accepted == "all" else 1] * len(gen_starts)) - _verify_step(b, SpeculativeRollback()) - _assert_state(b, _reference(b), accepted) - - def test_commit_without_staging_is_a_no_op(): ctx_lens, gen_starts, p, gen_accepted = CASES["decode"] b = _Batch(14, ctx_lens, gen_starts, p, _slots(len(gen_starts))) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_patch_embed_conv3d.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_patch_embed_conv3d.py index c8d4322ab1c0..21bdba2b0d0c 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_patch_embed_conv3d.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_patch_embed_conv3d.py @@ -5,7 +5,7 @@ The reference is a true-fp32 matmul of the flattened patch against the flattened filter: a different op, accumulation order and memory layout from the convolution under test, so it checks the contraction pairing rather than -restating it. +restating it. Beyond the cells: the one-patch-per-row guard and graph capture. """ import zlib @@ -13,7 +13,6 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import true_fp32_matmul from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.vision.patch_embed_conv3d import ( # noqa: E501 patch_embed_conv3d as op, ) @@ -89,125 +88,14 @@ def test_certified_cells(cell, certifying): # noqa: F811 _compare(out, op.reference(pixel_values, weight, bias)) -def test_five_dimensional_input_is_the_same_call() -> None: - pixel_values, weight, bias = _operands(64, seed=1) - flat = op(pixel_values, weight, bias) - explicit = op(pixel_values.view(-1, IN_CHANNELS, KT, KH, KW), weight, bias) - assert torch.equal(flat, explicit) - - -def test_zero_bias_is_the_identity() -> None: - pixel_values, weight, _ = _operands(128, seed=2, bias=False) - zero = torch.zeros(EMBED_DIM, dtype=weight.dtype, device="cuda") - assert torch.equal(op(pixel_values, weight, None), op(pixel_values, weight, zero)) - - -def test_rows_are_independent() -> None: - """Same shape, so cuDNN selects the same algorithm: asserted bitwise.""" - pixel_values, weight, bias = _operands(128, seed=5) - order = torch.randperm(128, generator=torch.Generator().manual_seed(0)).to("cuda") - out = op(pixel_values, weight, bias) - assert torch.equal(op(pixel_values[order].contiguous(), weight, bias), out[order]) - - -def test_wrong_variants_land_outside_the_gate() -> None: - """Three structural mis-derivations land hundreds of times outside - `_compare`'s atol; dropping the bias is bounded by `max|bias|` (~0.02 at the - tower's initializer scale) and is asserted to *be* the bias.""" - pixel_values, weight, bias = _operands(64, seed=6) - out = op(pixel_values, weight, bias).float() - ref = op.reference(pixel_values, weight, bias) - rms = ref.float().pow(2).mean().sqrt().item() - atol = max(1e-5, 4.0 * rms * 2.0**-8) - correct_dev = (out - ref.float()).abs().max().item() - rows = pixel_values.float() - - def _matmul(w): - with true_fp32_matmul(): - return rows @ w.reshape(w.shape[0], -1).float().t() + bias.float() - - structural = { - "row read as (T, C, Ph, Pw)": _matmul(weight.permute(0, 2, 1, 3, 4)), - "kernel flipped (true convolution)": _matmul(weight.flip(2, 3, 4)), - "patch rows and columns transposed": _matmul(weight.transpose(-1, -2)), - } - print(f"\n gate atol {atol:.5f}; correct deviation {correct_dev:.5f}") - for name, wrong in structural.items(): - dev = (out - wrong).abs().max().item() - print(f" {name:38s} dev {dev:9.5f} ({dev / atol:8.1f}x the gate)") - assert dev > 100.0 * atol, (name, dev, atol) - with pytest.raises(AssertionError): - _compare(out.to(ref.dtype), wrong.to(ref.dtype)) - - no_bias = (out - (ref.float() - bias.float())).abs().max().item() - print(f" {'bias dropped':38s} dev {no_bias:9.5f} ({no_bias / atol:8.1f}x the gate)") - assert no_bias > 5.0 * atol, (no_bias, atol) - assert abs(no_bias - bias.abs().max().item()) <= correct_dev - - -def test_guard_refuses_an_image_volume() -> None: - """The hazard, unguarded: a [1, C, D, H, W] volume with the element count of - 16 patches returns a [16, 1152] result of the right shape and dtype computed - from the wrong pixels.""" - g = torch.Generator(device="cuda").manual_seed(7) - volume = torch.empty(1, IN_CHANNELS, KT, 4 * KH, 4 * KW, device="cuda") - volume = volume.uniform_(-1.0, 1.0, generator=g).to(torch.bfloat16) - _, weight, bias = _operands(1, seed=7) - - reinterpreted = op(volume, weight, bias) - meant = ( - torch.nn.functional.conv3d(volume, weight, bias, stride=(KT, KH, KW)) - .flatten(2) - .squeeze(0) - .t() - ) - assert reinterpreted.shape == meant.shape == (16, EMBED_DIM) - dev = (reinterpreted.float() - meant.float()).abs().max().item() - print(f"\n unguarded volume reinterpretation deviates by {dev:.5f}") - assert dev > 0.1, dev - - with pytest.raises(AssertionError, match="each row must be exactly one patch"): - with validating(op): - op(volume, weight, bias) - - def test_guard_refuses_rows_that_are_not_one_patch() -> None: """Two patches per row are split into twice the rows, without an error.""" pixel_values, weight, bias = _operands(8, seed=8) doubled = pixel_values.reshape(4, 2 * PATCH_NUMEL) assert op(doubled, weight, bias).shape == (8, EMBED_DIM) - for bad in (doubled, pixel_values[:, :-1].contiguous(), pixel_values.reshape(-1)): - with pytest.raises(AssertionError, match="each row must be exactly one patch"): - with validating(op): - op(bad, weight, bias) - - -def test_dtype_and_device_mismatches_raise() -> None: - """Loud from torch itself, so the entry does not guard them.""" - pixel_values, weight, bias = _operands(8, seed=10) - with pytest.raises(RuntimeError, match="should be the same"): - op(pixel_values, weight.float(), bias.float()) - with pytest.raises(RuntimeError, match="should be the same"): - op(pixel_values, weight, bias.float()) - with pytest.raises(RuntimeError, match="should be the same"): - op(pixel_values, weight.cpu(), bias.cpu()) - - -def test_fp32_accumulates_in_tf32_unless_cudnn_is_told_otherwise() -> None: - pixel_values, weight, bias = _operands(256, dtype=torch.float32, seed=11) - ref = op.reference(pixel_values, weight, bias) - previous = torch.backends.cudnn.allow_tf32 - try: - torch.backends.cudnn.allow_tf32 = True - tf32_dev = (op(pixel_values, weight, bias) - ref).abs().max().item() - torch.backends.cudnn.allow_tf32 = False - exact = op(pixel_values, weight, bias) - finally: - torch.backends.cudnn.allow_tf32 = previous - exact_dev = (exact - ref).abs().max().item() - print(f"\n fp32 conv, cudnn TF32 on {tf32_dev:.3e}; off {exact_dev:.3e}") - _compare(exact, ref) - assert tf32_dev > 20.0 * exact_dev, (tf32_dev, exact_dev) + with pytest.raises(AssertionError, match="each row must be exactly one patch"): + with validating(op): + op(doubled, weight, bias) def test_cuda_graph_capture_and_replay() -> None: diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_pos_embed_interpolate.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_pos_embed_interpolate.py index deab3b492a53..3517e2f3d18e 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_pos_embed_interpolate.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_pos_embed_interpolate.py @@ -8,6 +8,7 @@ it gates the structure: which axis is which, how the table is read, and the spatial-merge row order. It sits up to 8.6e-6 from the exact value (its sample coordinates are float): far inside a bf16 cell's band, though not an fp32 one's. +Beyond the cells: the square-table guard and graph capture. """ import zlib @@ -45,15 +46,14 @@ def _merge_order(x, height, width, merge): ) -def _ref_resize(table, grid_thw, merge, mode="bilinear", align_corners=True): +def _ref_resize(table, grid_thw, merge): """Structural reference via torch's resize kernel, in fp64.""" num_grid = int(round(table.shape[0] ** 0.5)) hidden = table.shape[1] src = table.double().view(num_grid, num_grid, hidden).permute(2, 0, 1).unsqueeze(0).contiguous() outs = [] for frames, height, width in grid_thw: - kwargs = {} if mode == "nearest" else {"align_corners": align_corners} - resized = F.interpolate(src, size=(height, width), mode=mode, **kwargs) + resized = F.interpolate(src, size=(height, width), mode="bilinear", align_corners=True) rows = _merge_order(resized[0].permute(1, 2, 0), height, width, merge) outs.append(rows.repeat(frames, 1) if frames > 1 else rows) return torch.cat(outs) if len(outs) > 1 else outs[0] @@ -75,114 +75,6 @@ def test_certified_cells(cell, certifying): # noqa: F811 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.""" @@ -196,23 +88,6 @@ def test_guard_refuses_a_non_square_table() -> None: op(table, [(1, 4, 4)], MERGE) -def test_grid_not_divisible_by_the_merge_size_raises() -> None: - """Loud from torch, so not guarded.""" - table = _table(hidden=64, seed=11) - with pytest.raises(RuntimeError, match="is invalid for input of size"): - op(table, [(1, 5, 8)], MERGE) - with pytest.raises(RuntimeError, match="is invalid for input of size"): - op(table, [(1, 8, 5)], MERGE) - with pytest.raises(RuntimeError, match="is invalid for input of size"): - op(table, [(1, 2, 2)], 4) - - -def test_non_2d_table_raises() -> None: - table = torch.randn(NUM_POS, 4, 8, device="cuda", dtype=torch.bfloat16) - with pytest.raises(ValueError, match="weight has to be a 2D Tensor"): - op(table, [(1, 4, 4)], MERGE) - - def test_cuda_graph_capture_and_replay() -> None: """The grid is frozen at capture; the table's contents are not.""" table = _table(seed=12) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_vision_tower_calls.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_vision_tower_calls.py index 4096723e8342..d8a1f9396844 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_vision_tower_calls.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_vision_tower_calls.py @@ -124,20 +124,15 @@ def tower(): return M.VisionTower(_core()) -# `(grid, cap)`: frames of 256, 144 and 32 rows, of 32 and 8, and of 384. The cap cuts the -# mixed grid inside an item and between items, the one-item grid between frames of one video, -# and the short-frame grid into calls whose longest frame is at most 64 rows, where the fused -# attention kernel is the short-frame one. -CASES = { - "mixed": ([(2, 16, 16), (3, 12, 12), (6, 4, 8)], 512), - "short_frames_only": ([(9, 4, 8), (3, 2, 4)], 96), - "one_item_many_frames": ([(7, 16, 24)], 800), -} - - -@pytest.mark.parametrize("name", sorted(CASES)) -def test_calls_of_whole_frames_equal_one_call(tower, monkeypatch, name): - grid, cap = CASES[name] +# Frames of 256, 144 and 32 rows under a 512-row cap: three calls -- the first item, the +# second with two of the third's six frames, the other four -- so a cut between items and one +# inside an item; the last call's longest frame is 32 rows, where the fused attention kernel +# is the short-frame one. +GRID, CAP = [(2, 16, 16), (3, 12, 12), (6, 4, 8)], 512 + + +def test_calls_of_whole_frames_equal_one_call(tower, monkeypatch): + grid, cap = GRID, CAP pixels = _pixels(grid) grid_thw = torch.tensor(grid) monkeypatch.setattr(M, "_VISION_CALL_MAX_PATCHES", 1 << 20) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_kv_quant.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_kv_quant.py index 53e685592ac8..a699b909f1d7 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_kv_quant.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_kv_quant.py @@ -171,7 +171,7 @@ def _load(core, monkeypatch) -> None: core.post_load_weights() -def test_unquantized_kv_cache_keeps_the_checkpoints_bf16() -> None: +def test_unquantized_kv_cache_keeps_the_checkpoints_bf16(monkeypatch) -> None: core = _core(None) assert core.quant_mode == 0 assert core.kv_pool_dtype == torch.bfloat16 @@ -179,9 +179,14 @@ def test_unquantized_kv_cache_keeps_the_checkpoints_bf16() -> None: assert geometry.kv_dtype == torch.bfloat16 assert geometry.kv_width == core.head_dim assert core.kv_scale_orig_quant is None and core.kv_scale_quant_orig is None + _load(core, monkeypatch) + assert core.kv_scale_orig_quant is None and core.kv_scale_quant_orig is None -def test_fp8_kv_cache_stores_e4m3_under_the_fp8_quant_mode() -> None: +def test_fp8_kv_cache_stores_e4m3_with_unit_fp32_single_element_scales(monkeypatch) -> None: + """The generation kernels read `kvCacheScale[0]` and the QSA store and sparse GQA take + no scale, so the scale is a real tensor holding exactly 1.0: the dense and the sparse + path then agree on what a cache slot holds.""" core = _core(QuantAlgo.FP8) assert core.quant_mode == M.QM_FP8_KV_CACHE assert core.kv_pool_dtype == torch.float8_e4m3fn @@ -191,32 +196,6 @@ def test_fp8_kv_cache_stores_e4m3_under_the_fp8_quant_mode() -> None: assert geometry.kv_dtype == torch.float8_e4m3fn assert geometry.kv_width == core.head_dim assert geometry.dtype == torch.bfloat16 - - -def test_nvfp4_kv_cache_stores_e2m1_pairs_under_the_nvfp4_quant_mode() -> None: - core = _core(QuantAlgo.NVFP4) - assert core.quant_mode == M.QM_NVFP4_KV_CACHE - # An NVFP4 byte holds two values: the pool is uint8, half a head wide. - assert core.kv_pool_dtype == torch.uint8 - geometry = core._pool_geometry(32) - assert geometry.kv_dtype == torch.uint8 - assert geometry.kv_width == core.head_dim // 2 - assert geometry.dtype == torch.bfloat16 - - -@pytest.mark.parametrize("algo", [QuantAlgo.INT8, QuantAlgo.FP8_BLOCK_SCALES]) -def test_other_kv_cache_algorithms_are_refused(algo) -> None: - """The QSA backend serves its sparse path over half, bf16, fp8 and NVFP4 pools and falls - back to dense for any other, so a pool of another storage must not be accepted.""" - with pytest.raises(AssertionError, match="kv_cache_quant_algo"): - _core(algo) - - -def test_fp8_kv_cache_scales_are_unit_fp32_single_elements(monkeypatch) -> None: - """The generation kernels read `kvCacheScale[0]` and the QSA store and sparse GQA take - no scale, so the scale is a real tensor holding exactly 1.0: the dense and the sparse - path then agree on what a cache slot holds.""" - core = _core(QuantAlgo.FP8) _load(core, monkeypatch) for scale in (core.kv_scale_orig_quant, core.kv_scale_quant_orig): assert scale is not None @@ -224,11 +203,18 @@ def test_fp8_kv_cache_scales_are_unit_fp32_single_elements(monkeypatch) -> None: assert torch.equal(scale, torch.ones(1)) -def test_nvfp4_kv_cache_scales_are_unit_fp32_triples(monkeypatch) -> None: +def test_nvfp4_kv_cache_stores_e2m1_pairs_with_unit_fp32_scale_triples(monkeypatch) -> None: """The attention op reads the q, K and V second-level scales out of one fp32 `[3]` tensor, and the QSA NVFP4 store and sparse GQA take the same tensor, so every writer and reader applies the one value, 1.0 (the checkpoint has no KV scales).""" core = _core(QuantAlgo.NVFP4) + assert core.quant_mode == M.QM_NVFP4_KV_CACHE + # An NVFP4 byte holds two values: the pool is uint8, half a head wide. + assert core.kv_pool_dtype == torch.uint8 + geometry = core._pool_geometry(32) + assert geometry.kv_dtype == torch.uint8 + assert geometry.kv_width == core.head_dim // 2 + assert geometry.dtype == torch.bfloat16 _load(core, monkeypatch) for scale in (core.kv_scale_orig_quant, core.kv_scale_quant_orig): assert scale is not None @@ -236,7 +222,8 @@ def test_nvfp4_kv_cache_scales_are_unit_fp32_triples(monkeypatch) -> None: assert torch.equal(scale, torch.ones(3)) -def test_bf16_kv_cache_builds_no_scales(monkeypatch) -> None: - core = _core(None) - _load(core, monkeypatch) - assert core.kv_scale_orig_quant is None and core.kv_scale_quant_orig is None +def test_other_kv_cache_algorithms_are_refused() -> None: + """The QSA backend serves its sparse path over half, bf16, fp8 and NVFP4 pools and falls + back to dense for any other, so a pool of another storage must not be accepted.""" + with pytest.raises(AssertionError, match="kv_cache_quant_algo"): + _core(QuantAlgo.INT8) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_moe_scalars.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_moe_scalars.py index 35282a054b47..27a677f5931a 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_moe_scalars.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_moe_scalars.py @@ -47,11 +47,12 @@ def test_uniform_isc2_keeps_the_per_expert_products_bitwise() -> None: assert torch.equal(got["alpha1"], w["l0_e_isc1"] * w["l0_e_ws2_1"]) -@pytest.mark.parametrize("spread", [1.0001, 40.0]) -def test_isc2_spread_is_folded_into_alpha2(spread: float) -> None: - """Whatever the spread, `global_sf * alpha2[e]` is expert e's own `ws2_2[e]`, the - product trtllm-gen forms from `(1 / isc2[e]) * (isc2[e] * ws2_2[e])`.""" - isc2 = 0.01 * torch.linspace(1.0, spread, EXPERTS) +def test_isc2_spread_is_folded_into_alpha2() -> None: + """Whatever the spread (here 40x), `global_sf * alpha2[e]` is expert e's own `ws2_2[e]`, + the product trtllm-gen forms from `(1 / isc2[e]) * (isc2[e] * ws2_2[e])`: the FC2 value an + expert produces is the same whether its input was quantized at its own `1 / isc2[e]` or + at the layer's single scale.""" + isc2 = 0.01 * torch.linspace(1.0, 40.0, EXPERTS) w = _rows(isc2.flip(0)) cute = M._moe_layer_scalars(w, 0, "CUTEDSL") gen = M._moe_layer_scalars(w, 0, "TRTLLM") @@ -66,25 +67,6 @@ def test_isc2_spread_is_folded_into_alpha2(spread: float) -> None: assert torch.allclose(cute["global_sf"] * cute["alpha2"], w["l0_e_ws2_2"], rtol=1e-6) -def test_cutedsl_fc2_matches_the_per_expert_reference() -> None: - """The FC2 value an expert produces is `ws2_2[e] * (y @ W)` whether its input was - quantized at its own `1 / isc2[e]` (trtllm-gen) or at the layer's single scale.""" - isc2 = torch.tensor([0.02, 0.01, 0.04, 0.03, 0.02, 0.05, 0.01, 0.025]) - w = _rows(isc2) - cute = M._moe_layer_scalars(w, 0, "CUTEDSL") - gen = M._moe_layer_scalars(w, 0, "TRTLLM") - torch.manual_seed(0) - y = torch.randn(EXPERTS, 64, dtype=torch.float64) - weight = torch.randn(EXPERTS, 64, dtype=torch.float64) - ref = w["l0_e_ws2_2"].double() * (y * weight).sum(-1) - per_expert = gen["output2_scale_scalar"].double() * ( - (1.0 / isc2.double())[:, None] * y * weight - ).sum(-1) - one_scale = cute["alpha2"].double() * (cute["global_sf"].double() * y * weight).sum(-1) - torch.testing.assert_close(per_expert, ref, rtol=1e-5, atol=0) - torch.testing.assert_close(one_scale, ref, rtol=1e-5, atol=0) - - def test_gate_up_input_scale_mismatch_is_refused() -> None: w = _rows(torch.full((EXPERTS,), 0.0371)) w["l0_e_isc1_up"] = w["l0_e_isc1_up"] * 1.5 diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_multimodal_embeds.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_multimodal_embeds.py index 268f9a6ea448..5a371f24bf36 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_multimodal_embeds.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_multimodal_embeds.py @@ -111,17 +111,12 @@ def _item(modality, index): MIXED_ROWS = [200.0] * 4 + [100.0] * 4 + [201.0] * 8 + [101.0] -def test_image_request_is_unchanged(): +def test_the_splice_fills_the_placeholders_and_zeroes_the_padding(): + """An image request's rows land on its placeholders in one tower call and the padding + rows after the step are zeroed (`_splice` checks); a text step writes only the padding.""" tower = _Tower() - assert _splice(tower, _params(image=IMAGES), 5) == [100.0] * 4 + [101.0] + assert _splice(tower, _params(image=IMAGES), 5, padding=3) == [100.0] * 4 + [101.0] assert tower.calls == 1 - - -def test_the_splice_zeroes_the_padding_rows(): - assert _splice(_Tower(), _params(image=IMAGES), 5, padding=3) == [100.0] * 4 + [101.0] - - -def test_a_text_step_writes_only_the_padding(): target = SimpleNamespace(_step=SimpleNamespace(vision=None, num_tokens=2)) h = torch.full((4, HIDDEN), -1.0) M.PrefillTarget._splice(target, h) @@ -146,12 +141,6 @@ def test_a_mixed_request_needs_its_manifest(): _splice(_Tower(), _params(image=IMAGES, video=VIDEOS), len(MIXED_ROWS)) -def test_a_manifest_that_disagrees_with_the_payload_is_refused(): - params = _params(image=IMAGES, video=VIDEOS, order=MIXED_ORDER[:3]) - with pytest.raises(AssertionError, match="mm_item_order lists image items"): - _splice(_Tower(), params, len(MIXED_ROWS)) - - def test_a_chunk_takes_its_slice_of_the_prompt_order_rows(): runtime = SimpleNamespace( num_cached_mm_tokens=3, num_mm_tokens_in_chunk=6, total_embeds_in_request=len(MIXED_ROWS) @@ -261,18 +250,6 @@ def test_the_tower_calls_whole_frames_and_concatenates_their_rows_in_order(monke assert start == patches -def test_a_request_under_the_ceiling_is_one_call(monkeypatch): - grid = [(2, 4, 4), (1, 4, 8)] - tower = _tower(monkeypatch, 64) - pixels = _pixels(grid) - out = tower(pixels, torch.tensor(grid)) - assert len(tower.calls) == 1 - assert tower.calls[0][2].tolist() == [list(row) for row in grid] - # The one call sees the whole request and its output is the request's, row for row. - assert torch.equal(tower.calls[0][0], pixels) - assert torch.equal(out, pixels[:: MERGE * MERGE, :1].expand(-1, HIDDEN)) - - def test_a_request_past_the_old_whole_request_ceiling_is_served(monkeypatch): # 33 temporal groups of 32x32 patches (66 frames of 512x512 pixels): 33792 rows, 1024 # rows past what a single call may carry. @@ -285,14 +262,6 @@ def test_a_request_past_the_old_whole_request_ceiling_is_served(monkeypatch): assert out[:, 0].tolist() == [float(4 * j) for j in range(out.shape[0])] -def test_a_frame_past_the_ceiling_cannot_be_served(monkeypatch): - tower = _tower(monkeypatch, 64) - grid = [(1, 4, 4), (1, 10, 10)] - with pytest.raises(AssertionError, match="one frame is 10x10 = 100 patch rows"): - tower(_pixels(grid), torch.tensor(grid)) - assert tower.calls == [] - - def test_a_media_request_needs_the_processors_position_ids(): with pytest.raises(AssertionError, match="mrope_position_ids"): _splice(_Tower(), _params(video=VIDEOS, mrope=False), 12) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_pool_contract.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_pool_contract.py index 77909197a72d..20cdb53a3b25 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_pool_contract.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_pool_contract.py @@ -114,17 +114,9 @@ def test_valid_pools_pass() -> None: M.check_shared_selection(_t(2, 8, 512, dtype=torch.int32), {3: 1}, 3, 512) -@pytest.mark.parametrize( - "plane, match", - [ - (None, "did not allocate"), - (_t(9, 32, 3), "torch.bfloat16"), - (_t(9, 32, 2, dtype=torch.int32), "expected"), - ], -) -def test_position_plane_rejects(plane, match) -> None: - with pytest.raises(AssertionError, match=match): - M.check_position_plane(G, plane) +def test_position_plane_rejects_the_wrong_dtype() -> None: + with pytest.raises(AssertionError, match="torch.bfloat16"): + M.check_position_plane(G, _t(9, 32, 3)) def test_attention_pools_reject_wrong_index_plane_and_kv_layout() -> None: @@ -158,11 +150,7 @@ def test_attention_pools_reject_the_other_storage_dtype() -> None: # The index plane is a kernel contract, not KV storage: bf16 under either algorithm. with pytest.raises(AssertionError, match="QSA index plane"): M.check_attention_pools(G_FP8, 3, index.to(torch.float8_e4m3fn), fp8_kv) - - -def test_attention_pools_reject_the_wrong_nvfp4_width() -> None: - """An NVFP4 pool is half a head wide: a full-width uint8 pool is the other storage.""" - index = _t(9, 32, 1, 128) + # An NVFP4 pool is half a head wide: a full-width uint8 pool is the other storage. with pytest.raises(AssertionError, match="paged K/V"): M.check_attention_pools( G_NVFP4, 3, index, _t(9, 2, G.heads_kv, G.tokens_per_block, 256, dtype=torch.uint8) @@ -197,19 +185,10 @@ def test_kv_scale_planes_must_hold_the_data_pools_pages_per_slot() -> None: M.check_kv_scale_planes(G_NVFP4, 3, data, _strided((9, 2, 2, 32, 16), 24 * 1024 + 1)) -@pytest.mark.parametrize( - "scales, match", - [ - (None, "did not allocate"), - (_t(9, 2, 2, 32, 16), "torch.uint8"), - (_t(9, 2, 2, 32, 8, dtype=torch.uint8), "NVFP4 block scales"), - (_t(8, 2, 2, 32, 16, dtype=torch.uint8), "same slots"), - ], -) -def test_kv_scale_planes_reject(scales, match) -> None: +def test_kv_scale_planes_reject_the_wrong_block_scale_width() -> None: data = _t(9, 2, G.heads_kv, G.tokens_per_block, 128, dtype=torch.uint8) - with pytest.raises(AssertionError, match=match): - M.check_kv_scale_planes(G_NVFP4, 3, data, scales) + with pytest.raises(AssertionError, match="NVFP4 block scales"): + M.check_kv_scale_planes(G_NVFP4, 3, data, _t(9, 2, 2, 32, 8, dtype=torch.uint8)) def test_recurrent_pools_reject_wrong_key_dim_order_and_dtype() -> None: @@ -254,14 +233,21 @@ def test_step_buffers_reject_missing_conv_rank_dtype_and_index() -> None: M.check_step_buffers(G, int_conv, _conv(), None) -def test_legacy_state_rejects_missing_rank_and_dtype_drift() -> None: - int_ssm = _t(SLOTS, 2, G.num_v_heads, G.head_v_dim, G.head_k_dim, dtype=torch.float32) +def test_legacy_state_rejects_missing_rank_dtype_drift_and_too_many_rows() -> None: + def per_step(steps): + return _t(SLOTS, steps, G.num_v_heads, G.head_v_dim, G.head_k_dim, dtype=torch.float32) + + int_ssm = per_step(2) with pytest.raises(AssertionError, match="no per-step recurrent"): M.check_legacy_state(G, None, _ssm()) with pytest.raises(AssertionError, match="per-step recurrent buffer"): M.check_legacy_state(G, _ssm(), _ssm()) with pytest.raises(AssertionError, match="per-step recurrent buffer"): M.check_legacy_state(G, int_ssm.to(torch.bfloat16), _ssm()) + # More rows per request than the legacy recurrence is certified at. + M.check_legacy_state(G, per_step(M._MAX_LEGACY_VERIFY_WIDTH), _ssm()) + with pytest.raises(AssertionError, match="TRTLLM_USE_GDN_REPLAY"): + M.check_legacy_state(G, per_step(M._MAX_LEGACY_VERIFY_WIDTH + 1), _ssm()) def test_replay_state_rejects_pool_dtype_geometry_and_bookkeeping() -> None: @@ -281,15 +267,6 @@ def test_replay_state_rejects_pool_dtype_geometry_and_bookkeeping() -> None: M.check_replay_state(HISTORY, 3, pnat, buf.to(torch.int64), _ssm()) -def test_legacy_state_rejects_more_than_the_certified_rows_per_request() -> None: - def int_ssm(steps): - return _t(SLOTS, steps, G.num_v_heads, G.head_v_dim, G.head_k_dim, dtype=torch.float32) - - M.check_legacy_state(G, int_ssm(M._MAX_LEGACY_VERIFY_WIDTH), _ssm()) - with pytest.raises(AssertionError, match="TRTLLM_USE_GDN_REPLAY"): - M.check_legacy_state(G, int_ssm(M._MAX_LEGACY_VERIFY_WIDTH + 1), _ssm()) - - def test_verify_width_admits_one_to_eight_rows_per_request() -> None: for rows in (0, M._MAX_VERIFY_WIDTH + 1): with pytest.raises(AssertionError, match="predicted_tokens_per_seq"): diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_thop_attention_calls.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_thop_attention_calls.py index d9bbaa5c0ba0..87584e317806 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_thop_attention_calls.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_thop_attention_calls.py @@ -67,45 +67,33 @@ def _is_vision(call: ast.Call) -> bool: return "_VISION_CALL_CONSTANTS" in _unpacked(call) -def test_the_target_calls_thop_attention_where_expected() -> None: - """Guards the checks below against going vacuous: the trunk and the draft layer each - call it from the prefill and the decode target, and the vision tower once.""" - calls = _thop_attention_calls() - assert len([c for c in calls if not _is_vision(c)]) >= 4 - assert len([c for c in calls if _is_vision(c)]) == 1 - - _KV_CACHE_CALLS = [call for call in _thop_attention_calls() if not _is_vision(call)] -@pytest.mark.parametrize( - "call", _KV_CACHE_CALLS, ids=[f"call{i}" for i in range(len(_KV_CACHE_CALLS))] -) -def test_every_kv_cache_call_passes_the_pool_storage(call) -> None: - keywords = _keywords(call) - for name in _KV_STORAGE_ARGS: - assert name in keywords, f"thop_attention call at line {call.lineno} omits {name}" - value = keywords[name] - # Read off the core, whose `FP8` / `None` resolution all three follow. - assert isinstance(value, ast.Attribute) and value.attr == name, ( - f"line {call.lineno}: {name}={ast.unparse(value)}, not the core's `{name}`" - ) - - -def test_the_shared_constants_do_not_override_the_pool_storage() -> None: - """`**_CALL_CONSTANTS` is unpacked into every KV cache call; a storage argument in it - would either raise as a duplicate keyword or, if the call dropped its own, silently - pin the pool to one storage.""" +def test_every_kv_cache_call_passes_the_pool_storage() -> None: + # Against going vacuous: the trunk and the draft layer each call it from the prefill + # and the decode target. + assert len(_KV_CACHE_CALLS) >= 4 + for call in _KV_CACHE_CALLS: + keywords = _keywords(call) + for name in _KV_STORAGE_ARGS: + assert name in keywords, f"thop_attention call at line {call.lineno} omits {name}" + value = keywords[name] + # Read off the core, whose `FP8` / `None` resolution all three follow. + assert isinstance(value, ast.Attribute) and value.attr == name, ( + f"line {call.lineno}: {name}={ast.unparse(value)}, not the core's `{name}`" + ) + + +def test_the_unpacked_arguments_carry_no_pool_storage() -> None: + """`**_CALL_CONSTANTS` and `**step` (each target's `step_args`) are unpacked into every + KV cache call; a storage argument in either would raise as a duplicate keyword or, if the + call dropped its own, silently pin the pool to one storage.""" constants = {kw.arg for kw in _module_dict("_CALL_CONSTANTS").keywords} assert constants assert not constants & set(_KV_STORAGE_ARGS) for call in _KV_CACHE_CALLS: assert "_CALL_CONSTANTS" in _unpacked(call) - - -def test_the_step_state_carries_no_pool_storage() -> None: - """`**step` (each target's `step_args`) is unpacked into the same calls: a storage - argument there would collide with the call's own.""" steps = [ node for node in ast.walk(_TREE) @@ -121,11 +109,12 @@ def test_the_step_state_carries_no_pool_storage() -> None: def test_the_vision_tower_runs_without_kv_storage() -> None: - """The encoder-only call states a bf16, unscaled, cache-less configuration.""" + """The encoder-only call, made once, states a bf16, unscaled, cache-less configuration.""" constants = {kw.arg: kw.value for kw in _module_dict("_VISION_CALL_CONSTANTS").keywords} assert ast.literal_eval(constants["quant_mode"]) == 0 assert ast.literal_eval(constants["kv_scale_orig_quant"]) is None assert ast.literal_eval(constants["kv_scale_quant_orig"]) is None - for call in _thop_attention_calls(): - if _is_vision(call): - assert not set(_keywords(call)) & set(_KV_STORAGE_ARGS) + vision = [call for call in _thop_attention_calls() if _is_vision(call)] + assert len(vision) == 1 + for call in vision: + assert not set(_keywords(call)) & set(_KV_STORAGE_ARGS) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_validate_request.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_validate_request.py index b4ee457febf5..447ca2875a90 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_validate_request.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_validate_request.py @@ -86,6 +86,7 @@ def _item(modality, index): def test_clean_requests_are_admitted(): + # A single-modality request needs no manifest. _admit(_request(image=IMAGES)) _admit(_request(video=VIDEOS)) _admit(_request(image=IMAGES, video=VIDEOS, order=MIXED_ORDER)) @@ -105,11 +106,6 @@ def test_a_request_without_media_costs_nothing(): _admit(SimpleNamespace(py_multimodal_data={"multimodal_embedding": torch.zeros(4, 3)})) -def test_a_video_past_the_per_call_row_count_is_admitted(): - # 33 temporal groups of 32x32 patches: 33792 rows in all, 1024 past one call. - _admit(_request(video=[(33, 32, 32)])) - - def test_the_longest_admitted_frame_is_the_calls_ceiling(monkeypatch): monkeypatch.setattr(M, "_VISION_CALL_MAX_PATCHES", 64) _admit(_request(image=[(1, 8, 8)])) @@ -118,11 +114,10 @@ def test_the_longest_admitted_frame_is_the_calls_ceiling(monkeypatch): _admit(_request(image=[(1, 8, 8), (1, 8, 10)])) -@pytest.mark.parametrize("modality", ["image", "video"]) -def test_a_frame_past_the_ceiling_fails_the_request(modality): +def test_a_frame_past_the_ceiling_fails_the_request(): # 3072x3072 pixels: a 192x192 grid, 36864 patches in one frame. - request = _request(**{modality: [(1, 192, 192)]}) - with pytest.raises(ValueError, match=rf"{modality} 0 has frames of 192x192 = 36864 patches"): + request = _request(video=[(1, 192, 192)]) + with pytest.raises(ValueError, match=r"video 0 has frames of 192x192 = 36864 patches"): _admit(request) # The message names the pixels and the bound. with pytest.raises(ValueError, match=r"\(3072x3072 pixels\).*at most 32768 patches"): @@ -130,16 +125,17 @@ def test_a_frame_past_the_ceiling_fails_the_request(modality): def test_the_frame_bound_is_per_frame_not_per_item(): + # 33 temporal groups of 32x32 patches: 33792 rows in all, 1024 past one call. + _admit(_request(video=[(33, 32, 32)])) # 4 frames of 128x128 patches (16384 each) are fine; 192x192 in the second item is not. _admit(_request(video=[(4, 128, 128)])) with pytest.raises(ValueError, match="video 1 has frames of 192x192"): _admit(_request(video=[(4, 128, 128), (1, 192, 192)])) -@pytest.mark.parametrize("grid", [[(0, 4, 4)], [(1, 4, 5)]], ids=lambda g: str(g)) -def test_a_malformed_grid_is_refused(grid): +def test_a_malformed_grid_is_refused(): request = _request(image=[(1, 4, 4)]) - request.py_multimodal_data["image"]["image_grid_thw"] = torch.tensor(grid) + request.py_multimodal_data["image"]["image_grid_thw"] = torch.tensor([(1, 4, 5)]) with pytest.raises(ValueError, match="spatial_merge_size 2"): _admit(request) @@ -166,21 +162,16 @@ def test_a_payload_without_its_pixels_or_grid_is_refused(): _admit(request) -def test_missing_position_ids_fail_the_request(): +def test_missing_or_mis_sized_position_ids_fail_the_request(): request = _request(video=VIDEOS) + request.py_multimodal_data["mrope_config"]["mrope_position_ids"] = torch.zeros(3, 1, 40) + with pytest.raises(ValueError, match="one position per prompt token"): + _admit(request) del request.py_multimodal_data["mrope_config"] with pytest.raises(ValueError, match="mrope_position_ids.*got None"): _admit(request) -@pytest.mark.parametrize("shape", [(3, 1, 40), (2, 1, 12 + TEXT_TOKENS), (3, 12 + 5)]) -def test_mis_sized_position_ids_fail_the_request(shape): - request = _request(video=VIDEOS) - request.py_multimodal_data["mrope_config"]["mrope_position_ids"] = torch.zeros(*shape) - with pytest.raises(ValueError, match="one position per prompt token"): - _admit(request) - - def test_a_mixed_request_needs_its_manifest(): with pytest.raises(ValueError, match="needs mm_item_order"): _admit(_request(image=IMAGES, video=VIDEOS)) @@ -199,18 +190,6 @@ def test_a_manifest_that_disagrees_with_the_payload_is_refused(): _admit(_request(image=IMAGES, video=VIDEOS, order=[{"modality": "image"}])) -def test_a_single_modality_request_needs_no_manifest(): - _admit(_request(image=IMAGES, order=None)) - _admit(_request(video=VIDEOS, order=None)) - - -def test_a_prompt_whose_placeholders_disagree_with_the_media_is_refused(): - request = _request(video=VIDEOS) - request.py_multimodal_data["multimodal_embed_mask_cumsum"] = torch.arange(1, 9) - with pytest.raises(ValueError, match="holds 8 image/video placeholder tokens.*12 embedding"): - _admit(request) - - def test_the_placeholder_count_is_the_sum_over_both_modalities(): request = _request(image=IMAGES, video=VIDEOS, order=MIXED_ORDER) request.py_multimodal_data["multimodal_embed_mask_cumsum"] = torch.arange(1, 14) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_weight_table.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_weight_table.py index 7282046fbe35..155db4b45143 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_weight_table.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_weight_table.py @@ -86,12 +86,12 @@ def _core(mtp: bool, backend: str = "TRTLLM", prefill_copy: bool | None = None): ) -@pytest.mark.parametrize("mtp", [False, True], ids=["no_mtp", "mtp"]) -def test_every_declared_weight_has_a_source(mtp): +def test_every_declared_weight_has_a_source(): """A declared weight with no source allocates memory nothing ever fills -- silently zero rather than an error, which is what `load`'s coverage assert catches from the - other side.""" - core = _core(mtp) + other side. With MTP live: the rows without it are a subset + (`test_mtp_rows_appear_only_when_mtp_is_live`).""" + core = _core(True) manifest = W.MODEL_WEIGHTS.manifest(core) for key, sources in manifest.items(): assert sources, f"{key} is declared with no checkpoint source" @@ -260,13 +260,14 @@ def shared_expert(prefix): return w -@pytest.mark.parametrize("mtp", [False, True], ids=["no_mtp", "mtp"]) -def test_the_table_declares_what_the_init_block_did(mtp): - core = _core(mtp) +def test_the_table_declares_what_the_init_block_did(): + """With MTP live; without it the table drops exactly the MTP module's rows + (`test_mtp_rows_appear_only_when_mtp_is_live`).""" + core = _core(True) with torch.device("meta"): declared = W.MODEL_WEIGHTS.declare(W.MODEL_WEIGHTS.dims(core)) got = {k: (tuple(p.shape), p.dtype) for k, p in declared.items()} - want = _reference_declaration(mtp) + want = _reference_declaration(True) assert set(got) == set(want), ( sorted(set(got) - set(want))[:8], sorted(set(want) - set(got))[:8], @@ -368,9 +369,7 @@ def test_cutedsl_fc2_keeps_the_checkpoint_row_order_and_trtllm_gen_shuffles_it() assert int(shuffled[32 + 8 * v + u, 0]) == 32 + 4 * u + v -@pytest.mark.parametrize( - "backend, prefill_copy", [("TRTLLM", True), ("TRTLLM", False), ("CUTEDSL", False)] -) +@pytest.mark.parametrize("backend, prefill_copy", [("TRTLLM", True), ("CUTEDSL", False)]) def test_cute_dsl_expert_copy_rows_only_when_kept(backend, prefill_copy): """With the second expert copy (trtllm-gen, a step able to reach the CuteDSL dispatch threshold) every layer also carries its experts in the CuteDSL layout (`fc*_cd`); @@ -387,7 +386,6 @@ def test_cute_dsl_expert_copy_rows_only_when_kept(backend, prefill_copy): ("TRTLLM", 4096, None, True), ("TRTLLM", 4095, None, False), ("TRTLLM", 16384, "0", False), - ("TRTLLM", 16384, "1", True), ("CUTEDSL", 16384, None, False), ], ) From 8094d42e4e05f9e21e680c03d819268e2a6cda51 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Wed, 21 Oct 2026 03:19:14 -0700 Subject: [PATCH 27/28] [TRTLLM-17033][test] Qwen3.8-Flash-Next accuracy: acceptance floors for the fp8 and NVFP4 KV engines Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../references/acceptance_length.yaml | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/tests/integration/defs/accuracy/references/acceptance_length.yaml b/tests/integration/defs/accuracy/references/acceptance_length.yaml index 37918c1460d9..5f921108d166 100644 --- a/tests/integration/defs/accuracy/references/acceptance_length.yaml +++ b/tests/integration/defs/accuracy/references/acceptance_length.yaml @@ -157,12 +157,12 @@ ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp1: # test_mtp3_cutedsl_fp8_kv_bf16_state, over gsm8k: three drafts with the trunk's experts on # CuteDSL over an fp8 KV cache and a bf16 recurrent state. ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_cutedsl_fp8_kv_bf16_state: - # PENDING: modeling_v2's own full-gsm8k reading on GB300 (1319 samples); there is no built-in - # reading for this combination (the built-in model has no CuteDSL trunk). - ref_al: 3.705 + # modeling_v2's own full-gsm8k reading on GB300 (1319 samples); there is no built-in reading + # for this combination (the built-in model has no CuteDSL trunk). 3.705 over bf16 KV. + ref_al: 3.509 # 95% of ref_al, above the two-draft reading (~2.8): a third draft position that is never # accepted fails it. - min_al: 3.520 + min_al: 3.334 # test_mtp3_long_context_passkey: three drafts with IndexShare on, bf16 KV, over # 13 prompts of 3k-6k tokens (the sparse attention path). Most of the output is a copy of the @@ -180,17 +180,17 @@ ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_long_context: # test_mtp3_cutedsl_fp8_kv_bf16_state, over the pass keys of mtp3_long_context with IndexShare # off. ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_cutedsl_fp8_kv_bf16_state_long_context: - # PENDING - ref_al: 3.981 + # 3.962 here, against 4.000 over bf16 KV. + ref_al: 3.962 # 2.5% below ref_al, as for mtp3_long_context. - min_al: 3.88 + min_al: 3.86 # test_mtp7_long_context_passkey: seven drafts with IndexShare on, over an NVFP4 KV cache and a # bf16 recurrent state, over 12 prompts of 3k-12k tokens (the sparse attention path), whose -# output mostly copies the document's section headings: the readings sit near the range's -# ceiling of 8 (7.857 over bf16 KV, 7.774 over bf16 KV and an fp32 state). +# output mostly copies the document's section headings. The KV dtype sets the reading, and seven +# chained drafts amplify it: 7.857 over bf16 KV, 7.500 over fp8, 6.867 over NVFP4 (6.961 with +# IndexShare off, 7.003 over an fp32 state). ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp7_nvfp4_kv_bf16_state_long_context: - # PENDING - ref_al: 7.857 + ref_al: 6.867 # 2.5% below ref_al, as for mtp3_long_context. - min_al: 7.661 + min_al: 6.69 From 257d5f0ed475cc56c3cdf5007da3da3b4dcd974c Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Wed, 21 Oct 2026 19:37:21 -0700 Subject: [PATCH 28/28] [TRTLLM-17033][fix] PyExecutor admission: the model hook and its tests next to main's KV retention check _validate_request reached the model's optional hook through self.model_engine. test_connector_retention_request_validation (#20086) builds its executor with object.__new__(PyExecutor) and sets only the members the other checks read, so it failed with AttributeError on model_engine. The hook is now looked up through getattr(self, "model_engine", None), as _kv_connector_moves_kv reads kv_connector_manager: an executor without a model engine has no hook to call. #20086 also made _validate_request read _is_kv_manager_v2, which this PR's admission tests' executor stub lacked; _fetch_and_activate_new_requests turned the AttributeError into a rejection of every request. The stub now carries it. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 7 ++++--- tests/unittest/_torch/executor/test_py_executor.py | 1 + 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 5b361da5bea8..56c095e69cc1 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -6496,9 +6496,10 @@ def _validate_request(self, request: LlmRequest): # Perform sampler-specific validation self.sampler.validate_request(request) - # Perform model-specific validation, for models that define the hook - validate = getattr(getattr(self.model_engine, "model", None), - "validate_request", None) + # Perform model-specific validation, for models that define the hook. + # An executor without a model engine has no hook to call. + model = getattr(getattr(self, "model_engine", None), "model", None) + validate = getattr(model, "validate_request", None) if validate is not None: validate(request) diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 0f204af1154c..f4478bdb41ff 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -534,6 +534,7 @@ def _make_admission_executor(model, requests): real ones, over a model engine holding `model`; everything else they touch is a stub.""" executor = types.SimpleNamespace( model_engine=types.SimpleNamespace(model=model), + _is_kv_manager_v2=False, kv_cache_manager=None, kv_cache_transceiver=None, max_beam_width=1,